在深度学习图像分类任务中,我们常常面临一个现实困境:手头只有少量标注图片,却希望训练出一个性能不错的模型。从头训练一个深度神经网络需要海量数据和强大的算力,这对于个人开发者或小团队来说几乎不可能。幸运的是,迁移学习为我们提供了一条捷径。本文将手把手带你完成一个实战项目:利用 PyTorch 框架和预训练的 ResNet 模型,仅用少量图片,快速构建一个高效的图像分类器。无论你是刚入门深度学习的新手,还是希望快速验证想法的开发者,这套从环境搭建、数据准备、模型微调到结果评估的完整流程,都能让你直接复用,高效解决“小数据”下的分类难题。
1. 背景与核心概念:为什么需要迁移学习?
在深入代码之前,我们有必要厘清几个核心概念,理解迁移学习为何能成为小样本学习的“利器”。
深度学习与图像分类:深度学习,特别是卷积神经网络(CNN),在图像识别领域取得了革命性成功。它通过多层卷积、池化等操作,自动从像素中提取从边缘、纹理到物体部件的多层次特征,最终实现精准分类。
ResNet 模型:ResNet(残差网络)是 CNN 发展史上的一个里程碑。它通过引入“残差块”和“跳跃连接”,有效解决了深层网络中的梯度消失和网络退化问题,使得训练成百上千层的超深网络成为可能。在 ImageNet 等大型数据集上预训练的 ResNet 模型,已经学会了提取通用图像特征的强大能力。
迁移学习:这就是我们的核心武器。其核心思想是,将一个在源领域(如包含1000个类别的ImageNet数据集)大规模数据上训练好的模型(预训练模型)的知识,迁移到目标领域(我们自己的、数据量很少的特定分类任务)中来。我们不需要从零开始学习“如何看图片”,而是站在巨人的肩膀上,让模型快速适应新任务。
为什么它适用于少量图片的场景?
- 特征复用:预训练模型的前几层学习到的通常是通用特征(如边缘、角点、纹理),这些特征对于大多数视觉任务都是有用的。我们只需要让模型的后几层(负责高级语义和具体分类)去适应我们的新类别即可。
- 避免过拟合:数据量少时,训练参数众多的深层网络极易过拟合(模型只记住了训练集,而无法泛化到新数据)。迁移学习通过冻结大部分预训练参数(只更新少量新参数),极大地减少了需要学习的参数量,从而有效缓解过拟合。
- 加速收敛:模型已经有了一个非常好的权重初始化起点,因此训练过程收敛速度远快于从头训练。
接下来,我们将通过一个完整的实战案例,演示如何利用 PyTorch 和 ResNet 实现迁移学习。
2. 环境准备与版本说明
工欲善其事,必先利其器。首先确保你的开发环境已就绪。
操作系统:Windows 10/11, macOS 或 Linux (如 Ubuntu 20.04+) 均可。本文命令以 Linux/macOS 的 bash 和 Windows 的 PowerShell 为例。Python:推荐使用 Python 3.8 或 3.9。版本过高或过低可能导致某些库的依赖冲突。深度学习框架:我们将使用 PyTorch。请根据你的环境(是否支持CUDA)前往 PyTorch 官网 获取安装命令。
以下是一个在拥有 NVIDIA GPU 且已安装 CUDA 11.7 的环境下的安装示例:
# 使用 pip 安装 PyTorch、TorchVision 和 其他必要库 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 pip install matplotlib pandas tqdm pillow如果你的电脑没有 GPU,则安装 CPU 版本:
pip install torch torchvision torchaudio pip install matplotlib pandas tqdm pillow关键库版本参考(以写作时稳定版本为例,实际安装请以最新兼容版本为准):
torch: >= 1.13.0torchvision: >= 0.14.0matplotlib: >= 3.5.0 (用于可视化)pandas: >= 1.4.0 (可选,用于管理数据)tqdm: >= 4.64.0 (用于显示进度条)Pillow: >= 9.0.0 (图像处理)
IDE/编辑器:任意你熟悉的即可,如 VS Code、PyCharm 或 Jupyter Notebook。
项目结构:建议提前创建好如下目录结构,便于管理。
ai_study_transfer_learning/ ├── data/ # 数据目录 │ ├── train/ # 训练集 │ │ ├── class_a/ # 类别A的图片 │ │ ├── class_b/ # 类别B的图片 │ │ └── ... # 更多类别 │ └── val/ # 验证集(结构同train) ├── src/ # 源代码 │ ├── dataset.py # 自定义数据集类 │ ├── model.py # 模型定义 │ ├── train.py # 训练脚本 │ └── utils.py # 工具函数(如可视化) ├── outputs/ # 输出目录(保存模型、日志、图表) │ ├── models/ # 保存的模型权重 │ └── logs/ # 训练日志 └── requirements.txt # 项目依赖3. 核心原理与步骤拆解
在动手写代码前,让我们把迁移学习微调 ResNet 的流程和关键决策点梳理清楚。
3.1 迁移学习微调策略
对于 ResNet 这类预训练模型,我们通常采用以下两种微调策略:
特征提取器(Feature Extractor):
- 做法:冻结预训练模型的所有层(将其
requires_grad属性设为False),仅替换并训练最后的全连接分类层(fc层)。 - 适用场景:数据集非常小(每类几十张图),且与 ImageNet 数据集相似度较高。此时,预训练特征已经足够好,我们只需要学习一个简单的分类器。
- 优点:训练快,参数量少,极不易过拟合。
- 缺点:模型上限可能受限于预训练特征的通用性。
- 做法:冻结预训练模型的所有层(将其
微调全部层(Fine-tuning):
- 做法:不冻结任何层,替换最后的全连接层后,对整个模型的所有参数进行训练。通常,我们会使用一个较小的学习率,以免破坏预训练好的底层特征。
- 适用场景:数据集相对较大(每类几百张以上),或目标领域与 ImageNet 差异较大(如医学影像、卫星图片)。模型需要调整底层特征来更好地适应新数据。
- 优点:模型潜力更大,能更好地适应新任务。
- 缺点:训练慢,需要更多数据来防止过拟合,对学习率等超参数更敏感。
本文实战将采用一种更通用的折中策略:冻结大部分卷积层,微调最后几个块和全连接层。这是实践中最常用、效果最稳定的方法。
3.2 数据处理流程
由于数据量少,规范的数据处理至关重要,主要包括:
- 数据增强:对训练图像进行随机旋转、翻转、裁剪、色彩抖动等变换,在不增加真实数据的情况下,极大地扩充数据多样性,是防止过拟合的核心手段。
- 数据标准化:使用与预训练模型相同的均值和标准差对图像进行归一化。对于 TorchVision 中的 ResNet,通常使用 ImageNet 的统计量:
mean = [0.485, 0.456, 0.406],std = [0.229, 0.224, 0.225]。这能确保输入数据的分布与模型训练时一致。 - 数据集划分:必须划分独立的验证集(或测试集),用于在训练过程中监控模型泛化能力,防止过拟合。
3.3 模型修改关键点
以 ResNet-18 为例,我们需要修改其最后一层:
- 原始
fc层:Linear(in_features=512, out_features=1000)(对应 ImageNet 的1000类)。 - 新
fc层:Linear(in_features=512, out_features=N)(N 为我们自定义数据集的类别数)。
4. 完整实战案例:花卉分类
假设我们有一个小型的“花卉分类”数据集,包含daisy(雏菊)、rose(玫瑰)、sunflower(向日葵)三个类别,每个类别仅有 30-50 张训练图片。
4.1 准备数据集
请将你的图片按以下结构放置:
data/ ├── train/ │ ├── daisy/ │ │ ├── image1.jpg │ │ └── ... │ ├── rose/ │ │ ├── image1.jpg │ │ └── ... │ └── sunflower/ │ ├── image1.jpg │ └── ... └── val/ # 结构同train,每类约5-10张图 ├── daisy/ ├── rose/ └── sunflower/4.2 编写核心代码
4.2.1 定义数据集和数据加载器 (src/dataset.py)
我们使用torchvision.datasets.ImageFolder,它能自动根据文件夹结构生成标签。
# file: src/dataset.py import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader def get_data_loaders(data_dir, batch_size=32): """ 创建训练和验证数据加载器。 参数: data_dir: 数据根目录,包含'train'和'val'子目录。 batch_size: 批大小。 返回: train_loader, val_loader, class_names """ # 训练数据增强和归一化 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.RandomRotation(15), # 随机旋转±15度 transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动 transforms.ToTensor(), # 转换为Tensor [0,1] transforms.Normalize(mean=[0.485, 0.456, 0.406], # ImageNet统计量归一化 std=[0.229, 0.224, 0.225]) ]) # 验证/测试阶段,只进行中心裁剪和归一化,不做增强 val_transform = transforms.Compose([ transforms.Resize(256), # 将短边缩放到256 transforms.CenterCrop(224), # 中心裁剪224x224 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 使用ImageFolder加载数据集 train_dataset = datasets.ImageFolder(root=f'{data_dir}/train', transform=train_transform) val_dataset = datasets.ImageFolder(root=f'{data_dir}/val', transform=val_transform) # 创建数据加载器 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) # 获取类别名称 class_names = train_dataset.classes print(f"训练集样本数: {len(train_dataset)}") print(f"验证集样本数: {len(val_dataset)}") print(f"类别: {class_names}") return train_loader, val_loader, class_names if __name__ == '__main__': # 快速测试数据加载 train_loader, val_loader, classes = get_data_loaders('../data', batch_size=4) images, labels = next(iter(train_loader)) print(f"一个批次的图像形状: {images.shape}") # [4, 3, 224, 224] print(f"对应的标签: {labels}") # 如 tensor([0, 2, 1, 0])4.2.2 构建迁移学习模型 (src/model.py)
这是核心部分,我们加载预训练的 ResNet-18 并修改其最后一层。
# file: src/model.py import torch import torch.nn as nn from torchvision import models def create_model(num_classes, feature_extract=False, use_pretrained=True): """ 创建并配置用于迁移学习的ResNet-18模型。 参数: num_classes: 目标数据集的类别数量。 feature_extract: 如果为True,则冻结所有卷积层的参数,只训练全连接层。 use_pretrained: 是否使用在ImageNet上预训练的权重。 返回: model, input_size """ # 选择模型并加载预训练权重 model_ft = models.resnet18(weights=models.ResNet18_Weights.DEFAULT if use_pretrained else None) input_size = 224 # ResNet的标准输入尺寸 if feature_extract: # 特征提取模式:冻结所有网络参数 for param in model_ft.parameters(): param.requires_grad = False else: # 微调模式:我们选择冻结前面大部分层,只微调最后两个块 (layer3, layer4) 和 fc层 # 这样可以平衡训练速度和模型性能 for name, param in model_ft.named_parameters(): # 默认冻结所有参数 param.requires_grad = False # 解冻最后两个残差块和全连接层的参数 if 'layer3' in name or 'layer4' in name or 'fc' in name: param.requires_grad = True # 获取全连接层的输入特征数 num_ftrs = model_ft.fc.in_features # 替换最后的全连接层,使其输出维度等于我们的类别数 # 新的全连接层默认 requires_grad=True model_ft.fc = nn.Linear(num_ftrs, num_classes) # 打印可训练参数的数量 params_to_update = [] for name, param in model_ft.named_parameters(): if param.requires_grad: params_to_update.append(param) print(f"\t{name} 将被训练") print(f"模型总参数量: {sum(p.numel() for p in model_ft.parameters())}") print(f"可训练参数量: {sum(p.numel() for p in params_to_update)}") return model_ft, input_size if __name__ == '__main__': model, _ = create_model(num_classes=3, feature_extract=False) print(model)4.2.3 编写训练与验证脚本 (src/train.py)
# file: src/train.py import torch import torch.nn as nn import torch.optim as optim from torch.optim import lr_scheduler import time import copy from tqdm import tqdm import os def train_model(model, dataloaders, criterion, optimizer, scheduler, num_epochs=25, device='cpu', model_save_dir='../outputs/models'): """ 训练和验证模型。 参数: model: 要训练的模型。 dataloaders: 包含'train'和'val'键的字典,对应数据加载器。 criterion: 损失函数。 optimizer: 优化器。 scheduler: 学习率调度器。 num_epochs: 训练轮数。 device: 训练设备 ('cuda' 或 'cpu')。 model_save_dir: 模型保存目录。 返回: model: 训练好的最佳模型。 history: 包含训练历史的字典。 """ since = time.time() os.makedirs(model_save_dir, exist_ok=True) # 初始化最佳模型权重和准确率 best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []} for epoch in range(num_epochs): print(f'\nEpoch {epoch+1}/{num_epochs}') print('-' * 40) # 每个epoch都有训练和验证阶段 for phase in ['train', 'val']: if phase == 'train': model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss = 0.0 running_corrects = 0 # 使用tqdm显示进度条 pbar = tqdm(dataloaders[phase], desc=f'{phase.capitalize()} Epoch {epoch+1}') # 迭代数据 for inputs, labels in pbar: inputs = inputs.to(device) labels = labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 # 只在训练阶段追踪历史以计算梯度 with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 只在训练阶段进行反向传播和优化 if phase == 'train': loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) # 更新进度条描述 pbar.set_postfix({'Loss': loss.item()}) if phase == 'train' and scheduler is not None: scheduler.step() epoch_loss = running_loss / len(dataloaders[phase].dataset) epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset) # 记录历史 if phase == 'train': history['train_loss'].append(epoch_loss) history['train_acc'].append(epoch_acc.item()) else: history['val_loss'].append(epoch_loss) history['val_acc'].append(epoch_acc.item()) print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # 深度复制模型(如果验证准确率提高了) if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) # 保存最佳模型 best_model_path = os.path.join(model_save_dir, f'best_model_epoch{epoch+1}_acc{best_acc:.4f}.pth') torch.save({ 'epoch': epoch, 'model_state_dict': best_model_wts, 'optimizer_state_dict': optimizer.state_dict(), 'val_acc': best_acc, }, best_model_path) print(f'>>> 保存最佳模型至: {best_model_path}') # 每个epoch结束后保存一次检查点 checkpoint_path = os.path.join(model_save_dir, f'checkpoint_epoch{epoch+1}.pth') torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict() if scheduler else None, 'history': history, }, checkpoint_path) time_elapsed = time.time() - since print(f'\n训练完成,用时 {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s') print(f'最佳验证准确率: {best_acc:.4f}') # 加载最佳模型权重 model.load_state_dict(best_model_wts) return model, history def main(): # 设置设备 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print(f'使用设备: {device}') # 1. 加载数据 from dataset import get_data_loaders data_dir = '../data' # 修改为你的数据路径 batch_size = 16 # 根据GPU内存调整,数据量小时可以小一些 train_loader, val_loader, class_names = get_data_loaders(data_dir, batch_size=batch_size) dataloaders_dict = {'train': train_loader, 'val': val_loader} # 2. 创建模型 from model import create_model num_classes = len(class_names) model_ft, input_size = create_model(num_classes=num_classes, feature_extract=False, use_pretrained=True) model_ft = model_ft.to(device) # 3. 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() # 只优化那些 requires_grad=True 的参数 optimizer_ft = optim.SGD( filter(lambda p: p.requires_grad, model_ft.parameters()), lr=0.001, # 初始学习率,微调时通常设置较小 momentum=0.9 ) # 4. 定义学习率调度器:每7个epoch将学习率乘以0.1 exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) # 5. 训练模型 num_epochs = 15 # 对于小数据集,15-25个epoch通常足够 model_ft, hist = train_model( model_ft, dataloaders_dict, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=num_epochs, device=device ) # 6. 保存最终模型(可选) final_model_path = '../outputs/models/final_model.pth' torch.save(model_ft.state_dict(), final_model_path) print(f'最终模型已保存至: {final_model_path}') if __name__ == '__main__': main()4.2.4 可视化训练过程与模型预测 (src/utils.py)
# file: src/utils.py import matplotlib.pyplot as plt import torch import numpy as np from torchvision import transforms from PIL import Image def plot_training_history(history, num_epochs): """ 绘制训练和验证的损失、准确率曲线。 """ epochs = range(1, num_epochs + 1) plt.figure(figsize=(12, 4)) plt.subplot(1, 2, 1) plt.plot(epochs, history['train_loss'], 'b-', label='Training Loss') plt.plot(epochs, history['val_loss'], 'r-', label='Validation Loss') plt.title('Training and Validation Loss') plt.xlabel('Epochs') plt.ylabel('Loss') plt.legend() plt.grid(True) plt.subplot(1, 2, 2) plt.plot(epochs, history['train_acc'], 'b-', label='Training Accuracy') plt.plot(epochs, history['val_acc'], 'r-', label='Validation Accuracy') plt.title('Training and Validation Accuracy') plt.xlabel('Epochs') plt.ylabel('Accuracy') plt.legend() plt.grid(True) plt.tight_layout() plt.savefig('../outputs/training_history.png', dpi=150) plt.show() def predict_single_image(model, img_path, class_names, device='cpu', transform=None): """ 对单张图片进行预测。 """ # 加载图像 image = Image.open(img_path).convert('RGB') # 定义默认的验证集变换 if transform is None: transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 预处理 image_tensor = transform(image).unsqueeze(0) # 增加批次维度 image_tensor = image_tensor.to(device) # 预测 model.eval() with torch.no_grad(): outputs = model(image_tensor) _, predicted_idx = torch.max(outputs, 1) probabilities = torch.nn.functional.softmax(outputs, dim=1)[0] predicted_class = class_names[predicted_idx.item()] confidence = probabilities[predicted_idx.item()].item() # 显示结果 plt.figure(figsize=(8, 4)) plt.subplot(1, 2, 1) plt.imshow(image) plt.title(f'Predicted: {predicted_class}\nConfidence: {confidence:.2%}') plt.axis('off') plt.subplot(1, 2, 2) y_pos = np.arange(len(class_names)) plt.barh(y_pos, probabilities.cpu().numpy()) plt.yticks(y_pos, class_names) plt.xlabel('Probability') plt.title('Class Probabilities') plt.tight_layout() plt.show() return predicted_class, confidence if __name__ == '__main__': # 示例:加载历史并绘图 (假设历史数据已保存) # history = torch.load('../outputs/history.pth') # plot_training_history(history, num_epochs=15) pass4.3 运行与验证
- 准备数据:按照
4.1节的目录结构放置你的花卉图片。 - 安装依赖:在项目根目录创建
requirements.txt文件,内容参考2. 环境准备部分,然后运行pip install -r requirements.txt。 - 开始训练:在项目根目录下运行:
cd src python train.py - 监控训练:脚本运行时会打印每个 epoch 的损失和准确率,并自动保存最佳模型到
outputs/models/目录。 - 可视化结果:训练结束后,可以运行一个单独的脚本加载保存的
history和最佳模型,调用utils.py中的函数绘制学习曲线并对新图片进行预测。
4.4 预期结果
对于一个每类仅有几十张图片的小数据集,采用上述迁移学习方法,通常在 10-20 个 epoch 内,验证集准确率就能达到85% 甚至 90% 以上。这远高于从零开始训练(后者可能因过拟合而无法收敛)。你会观察到:
- 训练初期:验证准确率快速上升。
- 训练中后期:训练损失持续下降,但验证损失可能开始波动或上升,此时应关注验证准确率,并使用早停法或模型检查点来保存最佳模型,防止过拟合。
5. 常见问题与排查思路
在实际操作中,你可能会遇到以下问题:
| 问题现象 | 可能原因 | 解决思路 |
|---|---|---|
RuntimeError: CUDA out of memory | GPU 显存不足。 | 1. 减小batch_size(如从32降到16或8)。2. 使用更小的模型(如 ResNet-18 而非 ResNet-50)。 3. 在代码中使用 torch.cuda.empty_cache()清理缓存。4. 在 DataLoader中设置pin_memory=False。 |
| 训练损失不下降或准确率极低 | 1. 学习率设置不当。 2. 数据预处理错误(如归一化参数不对)。 3. 模型最后一层未正确修改。 4. 梯度被冻结( requires_grad=False)的层过多。 | 1. 尝试调整学习率(如 0.01, 0.001, 0.0001)。 2. 检查 transforms.Normalize的参数是否与预训练模型匹配。3. 打印模型结构,确认 fc层的out_features是否为你的类别数。4. 打印模型参数,确认需要训练的层是否已解冻。 |
| 验证准确率远低于训练准确率(过拟合) | 1. 训练数据太少。 2. 数据增强不够强。 3. 模型过于复杂或训练时间太长。 | 1. 尝试收集更多数据或使用更强的数据增强(如随机遮挡RandomErasing)。2. 增加 Dropout层(可在新的fc层前添加nn.Dropout(p=0.5))。3. 使用更早的检查点(验证准确率最高的模型)。 4. 尝试 feature_extract=True模式,进一步减少可训练参数。 |
KeyError: 'val_acc'或历史数据加载失败 | 保存的检查点文件结构与你加载时代码期望的结构不匹配。 | 确保加载模型权重的代码与保存时的字典键名一致。建议使用本文提供的标准化保存和加载方式。 |
| 预测时结果完全错误 | 1. 预测时未使用相同的预处理流程。 2. 图像通道顺序问题(OpenCV 读入为 BGR)。 3. 模型处于训练模式( model.train())。 | 1. 确保预测函数predict_single_image使用的transform与验证集完全相同。2. 如果使用 OpenCV ( cv2.imread),需将 BGR 转换为 RGB:image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)。3. 预测前务必调用 model.eval()。 |
6. 最佳实践与工程建议
掌握了基础流程后,以下建议能帮助你将项目提升到生产可用级别:
数据是王道:
- 质量高于数量:即使图片少,也要确保每张图片标注准确、清晰、有代表性。模糊、无关的图片有害无益。
- 严谨的划分:务必确保训练集和验证集没有重叠。对于小数据集,可以使用交叉验证来更可靠地评估模型性能。
- 高级增强:除了基本的翻转旋转,可以尝试
AutoAugment、RandAugment等策略,或使用albumentations库进行更丰富、针对性的增强(如透视变换、网格畸变等)。
模型选择与策略:
- 从小模型开始:数据量极少时,优先选择
ResNet-18甚至MobileNetV2、EfficientNet-B0等轻量模型,它们参数量少,更不易过拟合。 - 分层学习率:对模型不同部分使用不同的学习率。通常,新添加的
fc层使用较大的学习率(如 0.01),而微调的卷积层使用较小的学习率(如 0.001)。这可以通过optim.SGD的param_groups参数实现。 - 使用预训练权重:务必从官方或可靠来源加载预训练权重。TorchVision 的
weightsAPI 是最佳选择。
- 从小模型开始:数据量极少时,优先选择
训练过程监控与优化:
- 早停法:监控验证集损失,当其连续多个 epoch 不再下降时,提前停止训练,避免过拟合。
- 模型检查点:不仅保存最佳模型,也定期保存训练中间状态,便于回滚或分析。
- 使用 TensorBoard 或 WandB:替代简单的打印,使用可视化工具实时监控损失、准确率、权重分布等,效率更高。
- 超参数调优:对于关键超参数(如初始学习率、优化器类型、权重衰减系数),可以使用网格搜索或随机搜索,但要注意小数据集上的结果波动可能很大。
部署与推理优化:
- 模型导出:训练完成后,使用
torch.jit.script或torch.jit.trace将模型转换为 TorchScript,便于在非Python环境中部署。 - ONNX 格式:如果需要跨框架部署,可将模型导出为 ONNX 格式。
- 性能优化:对于生产环境,可以考虑使用
torch.quantization进行量化,或使用TensorRT、OpenVINO等推理加速引擎,显著提升推理速度。
- 模型导出:训练完成后,使用
代码与项目管理:
- 配置化:将数据路径、超参数、模型配置等写入
config.yaml文件,使代码更清晰,实验可复现。 - 版本控制:使用 Git 管理代码和实验记录。为每次重要的训练实验打上标签。
- 实验记录:记录每次实验的数据集、超参数、模型改动和最终结果,形成实验日志。
- 配置化:将数据路径、超参数、模型配置等写入
通过本次实战,你不仅学会了如何用少量图片和 ResNet 完成图像分类,更掌握了一套应对小样本深度学习任务的通用方法论。关键在于利用预训练模型的知识、通过数据增强创造“虚拟”数据、谨慎地微调网络参数,并始终以验证集为准绳防止过拟合。这套流程可以轻松迁移到其他视觉任务(如目标检测、图像分割)或其他预训练模型(如 Vision Transformer)。下一步,你可以尝试在自己的数据集上应用,探索不同的数据增强组合、调整微调层数、或者集成多个模型来进一步提升性能。