PyTorch迁移学习实战：用ResNet18实现20类食物图像分类（附代码详解）

一、迁移学习Transfer Learning详解1. 什么是迁移学习迁移学习是一种机器学习方法其核心思想是将从一个任务源任务中学到的知识应用到另一个相关但不同的任务目标任务中。它模仿了人类的学习方式——我们学习骑自行车后再学骑摩托车会容易很多学会识别猫狗后识别虎豹也会更轻松。在深度学习中迁移学习通常指在大规模数据集如ImageNet上预训练一个深度神经网络然后将该网络的权重作为初始值在目标任务的小规模数据集上进行微调或特征提取。2. 为什么需要迁移学习数据不足目标任务往往缺乏海量标注数据而从头训练大模型容易过拟合。计算资源有限从头训练ResNet等大型网络需要昂贵的GPU集群和时间成本。收敛更快预训练模型已经学会了通用的特征边缘、纹理、形状等只需在特定任务上微调即可快速达到良好性能。3. 迁移学习的常见策略策略做法适用场景特征提取冻结预训练模型的所有层不更新参数仅将输出特征送入新的分类器如全连接层进行训练。目标任务与源任务差异不大且目标数据量非常小。微调Fine-tuning解冻部分或全部预训练层允许参数在目标任务上继续更新。通常先用较小的学习率微调高层靠近输出的层再逐步微调低层。目标任务与源任务有一定差异且数据量中等几百到几千张。渐进式微调先冻结所有层训练新分类头再逐步解冻更多层联合训练。需要平衡训练速度和泛化能力。4. 关键注意事项数据分布差异若源任务ImageNet与目标任务如医学影像差异巨大迁移收益可能有限需考虑是否从零训练或使用域自适应技术。学习率选择微调时预训练层的学习率通常设为新层的1/10甚至更小避免破坏学到的通用特征。过拟合风险尽管迁移学习能缓解过拟合但若目标数据集仍然极小应使用更强的正则化Dropout、权重衰减、数据增强等。二、完整代码import torch import torchvision.models as models from torch import nn from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import numpy as np # 1. 加载预训练模型并冻结所有层 resnet_model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) for param in resnet_model.parameters(): param.requires_grad False # 2. 替换全连接层输出20类 in_features resnet_model.fc.in_features resnet_model.fc nn.Linear(in_featuresin_features, out_features20) # 3. 收集需要更新的参数仅新fc层 params_to_updata [] for param in resnet_model.parameters(): if param.requires_grad True: params_to_updata.append(param) # 4. 定义数据预处理 data_transforms { train: transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomRotation(10), transforms.CenterCrop(224), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomGrayscale(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), valid: transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) } # 5. 自定义Dataset类 class food_dataset(Dataset): def __init__(self, file_path, transformNone): self.file_path file_path self.imgs [] self.labels [] self.transform transform with open(self.file_path) as f: samples [x.strip().split( ) for x in f.readlines()] for img_path, label in samples: self.imgs.append(img_path) self.labels.append(label) def __len__(self): return len(self.imgs) def __getitem__(self, index): image Image.open(self.imgs[index]) if self.transform: image self.transform(image) label self.labels[index] label torch.tensor(int(label), dtypetorch.long) return image, label # 6. 加载数据集 train_dataset food_dataset(file_path./train.txt, transformdata_transforms[train]) test_dataset food_dataset(file_path./test.txt, transformdata_transforms[valid]) train_loader DataLoader(train_dataset, batch_size24, shuffleTrue) test_loader DataLoader(test_dataset, batch_size24, shuffleTrue) device cuda if torch.cuda.is_available() else cpu print(device) model resnet_model.to(device) loss_fn nn.CrossEntropyLoss() optimizer torch.optim.Adam(params_to_updata, lr0.001) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5) # 7. 训练函数 def train(dataloader, model, loss_fn, optimizer): model.train() for x, y in dataloader: x, y x.to(device), y.to(device) pred model(x) # 推荐使用 model(x) 而不是 model.forward(x) loss loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() # 8. 测试函数 best_acc 0 acc_s [] loss_s [] def test(dataloader, model, loss_fn): global best_acc size len(dataloader.dataset) model.eval() test_loss, correct 0, 0 with torch.no_grad(): for x, y in dataloader: x, y x.to(device), y.to(device) pred model(x) test_loss loss_fn(pred, y).item() correct (pred.argmax(1) y).sum().item() avg_loss test_loss / len(dataloader) accuracy correct / size print(fTest accuracy: {100*accuracy:.2f}%, avg loss: {avg_loss:.4f}) acc_s.append(accuracy) loss_s.append(avg_loss) if accuracy best_acc: best_acc accuracy torch.save(model.state_dict(), best_model.pth) # 保存最佳模型 # 9. 开始训练 epochs 10 for epoch in range(epochs): print(fEpoch {epoch1}/{epochs}) train(train_loader, model, loss_fn, optimizer) scheduler.step() test(test_loader, model, loss_fn) print(f最优训练结果准确率: {best_acc*100:.2f}%)三、核心知识点逐段解析3.1 冻结预训练模型for param in resnet_model.parameters(): param.requires_grad False将ResNet18的所有参数梯度设为False反向传播时不会计算这些参数的梯度也就不会更新它们。这样我们只训练新加的全连接层极大减少计算量。3.2 替换最后一层resnet_model.fc nn.Linear(in_features512, out_features20)ResNet18的fc层原是输出1000类ImageNet我们改成20类自己的食物类别。in_features可以通过resnet_model.fc.in_features动态获取。3.3 只收集可训练参数传给优化器params_to_updata [param for param in resnet_model.parameters() if param.requires_grad] optimizer torch.optim.Adam(params_to_updata, lr0.001)优化器只更新新fc层的参数提高效率。3.4 数据预处理与归一化transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])均值和标准差来自ImageNet数据集必须使用这些值才能与预训练模型匹配。3.5 自定义Dataset类__init__读取train.txt每行格式图片路径 标签将路径和标签分别存入列表。__getitem__根据索引加载图片、应用transform、返回(image, label)元组。3.6 训练/测试模式切换model.train() # 训练模式Dropout启用BN用batch统计量 model.eval() # 评估模式Dropout关闭BN用全局统计量3.7 学习率调度器scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5)每5个epoch将学习率乘以0.5有助于模型收敛更稳定。四、代码进一步优化建议当前代码已经能够正确运行并完成迁移学习任务以下优化建议可以让代码更健壮、更专业适合实际项目或进一步分享。4.1 保存最佳模型权重在测试函数中当发现更高的准确率时建议保存模型参数便于后续直接加载推理无需重新训练。if accuracy best_acc: best_acc accuracy torch.save(model.state_dict(), best_model.pth)4.2 移除冗余的 CenterCrop训练数据增强中的CenterCrop(224)在Resize(224,224)之后没有实际作用可以删除或替换为RandomResizedCrop(224)以获得更强的随机裁剪效果。transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224 transforms.RandomRotation(10), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomGrayscale(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])4.3 统一使用model(x)代替model.forward(x)model(x)是 PyTorch 推荐的调用方式它会自动调用forward并处理钩子等内部逻辑更规范。4.4 为数据集类增加convert(RGB)某些图片可能是灰度图L模式或RGBA模式为避免后续通道不匹配建议在打开图片后统一转为RGBimage Image.open(self.imgs[index]).convert(RGB)4.5 设置随机种子以保证结果可复现在脚本开头添加torch.manual_seed(42) np.random.seed(42)4.6 添加数据路径的鲁棒性检查在__init__中可以检查train.txt文件是否存在以及每行格式是否正确给出友好的错误提示。4.7 考虑使用验证集进行早停当前代码只记录了测试集的最佳准确率如果数据量允许建议将训练集再划分为训练验证集根据验证集的表现来调整超参数或提前停止训练避免在测试集上过拟合。五、实验结果与总结在20类食物数据集假设每类约100-200张上运行上述代码通常能获得75%~85%的测试准确率而从头训练相同模型可能只有40%~50%。这充分体现了迁移学习的威力。关键收获迁移学习是处理小数据集的利器可以大幅提升性能并节省时间。特征提取模式只需训练新分类头适合数据量很少的情况。数据预处理必须匹配预训练模型尺寸224×224ImageNet标准化。注意PyTorch中model.train()和model.eval()的正确使用否则会导致结果不稳定。细心处理标签转换和准确率计算避免低级bug。下一步学习方向尝试微调解冻最后几层卷积层用更低的学习率如1e-5继续训练。使用更先进的预训练模型ResNet50、EfficientNet、ViT等。学习学习率调整策略ReduceLROnPlateau、CosineAnnealing和早停技巧。希望这篇文章能帮助你彻底理解迁移学习的PyTorch实现。如果你有任何问题欢迎在评论区留言交流附完整代码已整合为可直接运行的脚本注意修改train.txt和test.txt的路径为你自己的数据集索引文件并修正路径中的逗号笔误。