PyTorch迁移学习实战:用ResNet预训练模型实现小样本图像分类

📅 2026/8/20 11:44:41
PyTorch迁移学习实战:用ResNet预训练模型实现小样本图像分类
1. 先搞清楚“少量图片”到底能少到什么程度做图像分类最头疼的就是数据不够。标注几百张、几千张图片对很多实际项目来说成本太高。这时候迁移学习就成了最直接的解决方案。它解决的核心问题是如何利用别人在大数据集如ImageNet上训练好的强大模型来快速适应你自己的、数据量很少的新任务。这次我们实战的主题就是用PyTorch基于ResNet预训练模型用很少的图片来训练一个图像分类器。这里的“少量”不是理论概念而是有明确边界的。根据我的经验每个类别能有50到200张质量不错的图片就足够让模型学到有效的特征达到不错的实用精度。如果每个类别只有10张图那就要在数据增强上花更多心思并且对最终效果的预期要更现实一些。这个方法特别适合个人开发者或小团队没有海量标注预算想快速验证一个图像分类的想法。特定垂直领域比如识别特定种类的工业零件、稀有动植物、医疗影像中的特定征象这些领域公开数据少但预训练模型的特征提取能力通用。学习入门想跳过从零训练CNN的漫长过程直接感受模型调优和解决实际问题的流程。整个流程的关键不是从头设计网络而是做好两件事一是正确地“改造”预训练模型让它适应你的新任务比如从1000类ImageNet分类改成你的5类分类二是用一套“组合拳”防止在小数据集上过拟合。下面我们就一步步拆解。2. 环境与数据准备别在第一步踩坑在写任何代码之前先把环境和数据整理清楚能避免一大半后续的诡异报错。2.1 环境配置要点你需要一个Python环境3.7以上均可和PyTorch。对于图像任务GPU不是必须的但能极大缩短训练时间。用CPU也能跑只是等待时间会长很多。# 基础安装命令示例请根据你的CUDA版本选择 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # CUDA 11.8 # 或者CPU版本 pip install torch torchvision torchaudio我建议直接使用torchvision它内置了ResNet等经典模型的预训练权重和标准数据加载方式省去很多麻烦。另外可以安装matplotlib和tqdm用于可视化和进度显示。关键检查点PyTorch版本不同版本的torchvision对模型加载可能有细微差别。建议保持较新稳定版如2.0。CUDA/cuDNN如使用GPU确保torch.cuda.is_available()返回True。磁盘空间预训练模型文件如ResNet50大约100MB下载前确保有空间。2.2 数据组织是重中之重数据怎么放直接决定了后面代码怎么写。请严格按照这个结构来组织你的图片你的数据集根目录/ ├── train/ │ ├── class_1/ │ │ ├── img_001.jpg │ │ ├── img_002.jpg │ │ └── ... │ ├── class_2/ │ │ ├── img_003.jpg │ │ └── ... │ └── ... └── val/ ├── class_1/ │ ├── img_101.jpg │ └── ... ├── class_2/ │ └── ... └── ...train和val或test目录是必须的代表训练集和验证集。每个子目录的名称就是类别标签例如cat,dog,car。torchvision.datasets.ImageFolder会自动根据子目录名来分配标签。即使图片很少也务必划分验证集这是判断模型是否过拟合的唯一可靠依据。通常可以按8:2或7:3的比例划分训练集和验证集。图片格式常见jpg、png均可但注意统一。如果图片尺寸不一没关系我们后面会用transforms统一调整。针对“少量图片”的数据增强策略 因为数据少我们需要在代码里对训练集图片进行实时数据增强来“凭空”创造更多的训练样本变体。这是提升小数据集泛化能力的关键。常见的增强包括随机翻转、旋转、裁剪、颜色抖动等。验证集则不能做随机增强通常只做尺寸调整和归一化保证评估的稳定性。3. 模型加载与改造理解“冻结”与“微调”这是迁移学习的核心步骤。我们不是从随机初始化的权重开始训练而是从一个在ImageNet上已经训练好的、能提取通用图像特征如边缘、纹理、形状的ResNet模型开始。3.1 加载预训练模型import torchvision.models as models import torch.nn as nn # 以ResNet50为例pretrainedTrue会下载预训练权重 model models.resnet50(pretrainedTrue)此时model的最后一层全连接层通常是fc或classifier是针对ImageNet的1000个类别的。我们需要替换它。3.2 替换分类头关键操作假设我们的新任务有5个类别。num_ftrs model.fc.in_features # 获取原全连接层的输入特征数 model.fc nn.Linear(num_ftrs, 5) # 替换为一个新的全连接层输出为5现在模型的特征提取部分卷积层保留了预训练权重而新的分类头是随机初始化的。3.3 决定冻结哪些层策略选择这是影响训练速度和效果的关键决策。有两种主流策略特征提取器模式冻结卷积层 将除了最后全连接层之外的所有层的权重“冻结”requires_grad False不让它们在训练中更新。只训练新替换的fc层。for param in model.parameters(): param.requires_grad False # 只让最后一层可训练 for param in model.fc.parameters(): param.requires_grad True适用场景数据量非常少每类50或者你只想快速得到一个基线模型。训练速度快但性能上限可能较低。微调模式Fine-tuning 解冻所有层或解冻靠近输出的部分卷积层例如ResNet的layer4让它们和新分类头一起参与训练。# 解冻所有层默认就是True如果之前冻结过需要再打开 for param in model.parameters(): param.requires_grad True # 或者更精细地只解冻最后几个块 # for name, param in model.named_parameters(): # if layer4 in name or fc in name: # param.requires_grad True # else: # param.requires_grad False适用场景数据量相对多一些每类100且新任务的数据分布与ImageNet自然图像有一定差异。训练速度慢需要更小的学习率但通常能达到更好的效果。我的建议对于新手可以先从特征提取器模式开始跑通流程并观察验证集精度。如果精度不满足要求再尝试切换到微调模式并使用更小的学习率例如比特征提取时小10倍。4. 训练流程拆解数据、损失、优化与迭代模型准备好后就要进入训练循环。这里每一步都有需要注意的细节。4.1 数据加载与增强使用torchvision的ImageFolder和DataLoader。import torchvision.transforms as transforms from torchvision import datasets from torch.utils.data import DataLoader # 定义训练和验证的数据增强与变换 data_transforms { train: transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪到224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ToTensor(), # 转为Tensor transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 归一化 ]), val: transforms.Compose([ transforms.Resize(256), # 先缩放到256 transforms.CenterCrop(224), # 再从中心裁剪224 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 加载数据集 image_datasets {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [train, val]} # 创建数据加载器 dataloaders {x: DataLoader(image_datasets[x], batch_size32, shuffleTrue if xtrain else False, num_workers4) for x in [train, val]}参数解释与避坑224这是ResNet等模型的标准输入尺寸。不要随意更改除非你清楚网络结构并能对应修改。Normalize的均值标准差这是ImageNet数据集上的统计值。使用预训练模型时必须用同样的值归一化你的输入否则特征分布不对效果会大打折扣。batch_size根据你的GPU显存调整。小数据下4, 8, 16, 32都常见。太小可能导致训练不稳定。num_workers数据加载的并行进程数可以加快数据读取。在Windows上有时设为0更稳定。4.2 损失函数与优化器import torch.optim as optim criterion nn.CrossEntropyLoss() # 多分类任务的标准损失函数 # 优化器只训练需要梯度的参数 optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.001, momentum0.9) # 或者使用Adam # optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.001)学习率设置是关键如果采用特征提取器模式只训练最后一层学习率可以设得稍大如0.01, 0.001。如果采用微调模式训练很多层学习率必须调小如0.0001, 0.00001否则预训练好的权重会被快速破坏导致模型“失忆”。更高级的做法是使用学习率调度器如StepLR或ReduceLROnPlateau在训练过程中动态降低学习率。4.3 训练与验证循环这是标准的PyTorch训练模板但有几个地方需要特别关注。def train_model(model, criterion, optimizer, num_epochs25): for epoch in range(num_epochs): print(fEpoch {epoch}/{num_epochs - 1}) print(- * 10) # 每个epoch都有训练和验证阶段 for phase in [train, val]: if phase train: model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss 0.0 running_corrects 0 # 迭代数据 for inputs, labels in dataloaders[phase]: 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) epoch_loss running_loss / dataset_sizes[phase] epoch_acc running_corrects.double() / dataset_sizes[phase] print(f{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}) # 这里可以添加模型保存逻辑例如保存验证集上最好的模型 # if phase val and epoch_acc best_acc: # best_acc epoch_acc # torch.save(model.state_dict(), best_model.pth) return model需要盯住的点model.train()和model.eval()必须正确切换。训练模式会启用Dropout、BatchNorm的更新等评估模式会关闭它们保证结果稳定。optimizer.zero_grad()每次计算新梯度前要清零否则梯度会累积。with torch.set_grad_enabled(phase train)在验证阶段我们不需要计算梯度这个上下文管理器可以节省内存和计算。损失和准确率重点看验证集val的指标。训练集损失一直降、准确率一直升是正常的但如果验证集指标在几轮后不再提升甚至下降说明过拟合了。5. 结果分析与问题排查你的模型真的学好了吗训练完成后不要只看最后的准确率数字。有几个步骤能帮你判断模型质量并定位问题。5.1 可视化学习曲线将每个epoch的训练损失/准确率和验证损失/准确率画出来这是最直观的诊断工具。理想情况训练和验证损失同步下降准确率同步上升最后趋于平稳。过拟合训练损失持续下降、准确率持续上升但验证损失在某个点后开始上升验证准确率停滞或下降。这说明模型只记住了训练数据没学会泛化。欠拟合训练和验证损失都下降得很慢准确率都很低。说明模型能力不足或训练不够可能学习率太小、epoch太少、或者模型被过度冻结。5.2 在验证集上查看混淆矩阵准确率只是一个总和。你需要知道模型具体在哪些类别上分不清。from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns # ... 在验证集上运行模型收集所有预测和真实标签 ... all_preds [] all_labels [] with torch.no_grad(): for inputs, labels in dataloaders[val]: inputs inputs.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) # 计算混淆矩阵 cm confusion_matrix(all_labels, all_preds) # 可以打印分类报告查看精确率、召回率等 print(classification_report(all_labels, all_preds, target_namesclass_names))混淆矩阵能清晰显示有多少“猫”被误判为“狗”。如果某些类别总是混淆可能意味着这些类别的图片本身视觉上就很相似。这些类别的训练样本数量严重不平衡。数据增强不够模型没有看到足够的视角变化。5.3 常见问题与排查清单如果效果不理想按这个顺序检查数据问题最常见图片路径或标签错误用几行代码打开几张图片确认ImageFolder加载的标签和图片是对应的。数据泄露训练集和验证集的图片有重复。确保划分时是随机的、无重叠的。类别极度不平衡某个类别的图片特别少。考虑使用加权的损失函数nn.CrossEntropyLoss(weightclass_weights)或对该类进行过采样。图片质量或标注质量差图片模糊、标注错误。没有好的数据再好的模型也没用。预处理问题归一化参数用错确认使用了[0.485, 0.456, 0.406], [0.229, 0.224, 0.225]。输入尺寸不对不是224x224。检查transforms的设置。模型与训练问题学习率不合适这是微调时最大的坑。如果微调先把学习率降到0.0001或更小试试。没有冻结/错误冻结想用特征提取模式却忘了冻结层导致所有层都在大学习率下乱跑。检查requires_grad。过拟合训练集准确率远高于验证集。解决方案加强数据增强、添加Dropout、减少模型复杂度用更小的ResNet如ResNet18、使用早停Early Stopping、或者收集更多数据。欠拟合训练集准确率就很低。解决方案减少冻结层解冻更多层进行微调、增加训练轮数、适当增大学习率、检查数据是否有问题。代码细节问题忘记model.train()和model.eval()导致训练时BatchNorm用全局统计或验证时用了Dropout。GPU/CPU不一致模型、数据、标签要在同一个设备上model.to(device),data.to(device)。6. 进阶优化与生产化思考当你的模型在验证集上表现稳定后可以考虑以下步骤让它从一个实验脚本变得更“可用”。6.1 模型保存与加载保存最好的模型通常是在验证集上性能最好的那个而不是最后一个epoch的模型。# 在验证循环内 if phase val and epoch_acc best_acc: best_acc epoch_acc # 保存整个模型结构权重 torch.save(model, best_model_full.pth) # 或者只保存权重推荐更灵活 torch.save(model.state_dict(), best_model_weights.pth) # 加载时 # 方式一加载完整模型需要模型类定义在作用域内 model torch.load(best_model_full.pth) model.eval() # 方式二加载权重到定义好的模型结构中 model.load_state_dict(torch.load(best_model_weights.pth)) model.eval()6.2 编写预测函数训练好的模型最终是用来做预测的。def predict_image(image_path, model, transform, class_names): model.eval() image Image.open(image_path).convert(RGB) image_tensor transform(image).unsqueeze(0) # 增加batch维度 image_tensor image_tensor.to(device) with torch.no_grad(): outputs model(image_tensor) _, predicted_idx torch.max(outputs, 1) probability torch.nn.functional.softmax(outputs[0], dim0) predicted_label class_names[predicted_idx.item()] confidence probability[predicted_idx].item() return predicted_label, confidence注意预测时使用的transform必须和验证集的一模一样通常是val的变换不做随机增强。6.3 尝试不同的BackboneResNet50不错但未必是最优解。对于小数据任务有时更小、更快的模型反而更好因为它们参数少更不容易过拟合。ResNet18速度快参数量少在小数据上表现往往不差是首选的基线模型。EfficientNet在准确率和效率之间平衡得更好但加载预训练权重时要注意版本匹配。MobileNet为移动端设计速度极快精度略有牺牲。更换Backbone通常只需要修改一行加载模型的代码但要注意不同模型最后的分类层名称可能不同有的是fc有的是classifier。6.4 考虑更系统的实验管理如果任务很重要你需要更严谨地实验使用TensorBoard或Weights Biases可视化损失曲线、准确率、甚至输入图片和模型权重分布。交叉验证对于极小的数据集可以使用交叉验证来更可靠地评估模型性能。超参数调优系统性地调整学习率、优化器、数据增强强度、冻结层数等。可以手动尝试也可以使用如Optuna等自动化工具。用ResNet做迁移学习进行小样本图像分类本质上是一个“站在巨人肩膀上”的工程实践。成功的关键不在于理论多深奥而在于对数据、模型、训练流程这些基础环节的扎实理解和细致操作。先从特征提取模式跑通一个基线记录结果然后尝试微调对比效果分析错误迭代优化。这个过程本身就是深度学习项目落地最核心的经验。