PyTorch迁移学习实战：14类水稻虫害图像分类完整流程

水稻虫害分类这个项目听起来像是典型的“用深度学习做个图像分类”任务但真正动手做的时候数据整理、预处理、迁移学习微调、评估指标这些环节每一个都会卡人。这篇文章就用PyTorch和torchvision这套组合以14类水稻虫害数据集为例把从数据准备到模型训练、评估的完整流程拆开讲清楚包含可以直接抄走的代码和踩坑经验。无论你是刚接触深度学习的学生还是在农业AI项目里被数据折腾过的工程师这套路线都能让你少走不少弯路。1. 整体思路与方案选型拆解1.1 为什么选择预训练ResNet模型接手一个14类水稻虫害分类任务我脑子里第一反应不是从头设计网络而是去确认能不能用现成的预训练模型。原因很现实田间拍摄的水稻害虫照片受光照、背景、叶片遮挡、拍摄距离的影响非常大同一个虫害在不同环境下外观差异明显分类任务本质上是“在通用视觉特征之上做细粒度区分”。ResNet系列的残差结构让网络在加深后依然能稳定收敛torchvision里直接提供resnet18、resnet34、resnet50这些主力型号并且自带ImageNet预训练权重。这套权重在通用特征上已经具备很强的表达能力边缘、纹理、颜色模式这些底层视觉元素都能直接复用。我的实际经验是农业虫害数据集很少有大几十万张的量级每个类能有几百到几千张图就算不错了。这时候如果从零初始化一个深层网络训练几百个epoch都不一定收敛得动泛化能力也差。预训练模型正好同时解决数据量不足和训练时间长的双重问题。具体怎么用分两种思路一种是冻住backbone所有参数只训练最后的分类头另一种是加载预训练权重后让整个网络一起以小学习率微调。对14类水稻虫害来说我的结论是绝大多数情况用全量微调效果更好。因为水稻稻飞虱、二化螟、稻纵卷叶螟这类害虫间的差异往往体现在翅纹、斑点和体型的细粒度特征上光靠ImageNet迁移来的通用特征还不够必须让网络的高层也跟着适应农业图像的统计分布。结合ResNet50来说最后全连接层之前的特征输出是2048维改成14分类只需要把fc层替换成Linear(2048, 14)这一层几乎不用动其他代码。工程上ResNet也不是最轻量级的模型但在单张主流显卡上batch_size调到32甚至64都没有压力农业项目里我们优先考虑精度而不是极限推理速度。1.2 为什么选择PyTorch torchvisionPyTorch在迁移学习场景中的优势太明显了。torchvision.models里一行代码就能把网络结构和预训练权重一起载入不再需要手动下载权重文件再手工填进状态字典。预处理这块torchvision.transforms直接提供标准化工具用ImageNet预训练权重时归一化的均值和标准差就是固定的mean[0.485, 0.456, 0.406]、std[0.229, 0.224, 0.225]这个细节极其关键——图像输入模型的数值范围不对效果会差出一大截很多人一开始忘了做归一化模型精度掉得莫名其妙都不知道去哪查。还有一点PyTorch的DataLoader、Dataset接口写起来跟写普通Python类没有本质区别。动态图特性让调试变得直观遇到维度不匹配、梯度异常能很快定位。选择torchvision的隐性好处是它的模型定义本身就是标准参考实现。当你要修改输出层、冻结参数、抽取中间层特征时网络各模块的名称是确定的不会出现包装太深导致无法操作的情况。配合预训练权重核心代码量控制在几百行以内这对工程落地和后续维护都很有价值。1.3 从数据到模型的完整技术路线结合这次14类水稻虫害分类任务我习惯按照下面这条主线推进确认原始数据形态整理成train/val/test三类目录每个类别单独一个子目录用ImageFolder或自定义Dataset构建训练集和验证集的读取逻辑在transforms里给训练集加入随机裁剪、翻转、颜色抖动等增强验证集只做固定尺寸和归一化加载预训练ResNet替换最后全连接层决定是否冻结主干配置损失函数、优化器、学习率调度器把训练循环封装成函数每个epoch结束后在验证集上评估保存验证集表现最好的权重训练完成后加载最优权重在测试集上输出分类报告和混淆矩阵等指标。这个流程把“一定会用到的环节”和“需要按实际情况调整的环节”分开了。数据增强策略可以根据类别分布与图片质量调整冻结主干与否可以根据预训练模型与数据的相似程度调整但整体框架是稳定的。下面就是把这条路线一步一步落地的实际操作。2. 数据集准备与预处理细节2.1 数据集结构与14类标签规划先说最常见的场景。甲方给到的数据通常是压缩包解压后一堆文件夹每类一个名字或者是一个大文件夹里图片文件名带类别前缀。最省事的整理方式是做成PyTorch的ImageFolder标准结构dataset/ train/ brown_planthopper/ img_001.jpg img_002.jpg rice_leaf_roller/ img_003.jpg ... val/ brown_planthopper/ rice_leaf_roller/ ... test/ ...train和val下面每个子文件夹的名字就是类别名。14个类别我建议先把类别顺序记录在一个文本文件里因为之后模型输出的索引与类别名的对应关系必须保持一致。比如按字母排序后索引0对应brown_planthopper索引1对应rice_stem_borer这个映射在写混淆矩阵、做推理脚本时都绕不开。如果原始数据是散乱图片加一个Excel标签表那就需要先把标签表读出来按行把图片移动到对应类别文件夹或者写一个自定义Dataset直接读路径列表和标签数组。我更推荐前者因为ImageFolder稳定成熟后续做推理脚本也省事。这里有个特别提醒田间采集的数据经常是jpg、png、bmp混在一起偶尔还会出现损坏的图片文件。我一般先写几行代码把所有图片逐张用PIL打开并verify一遍把损坏文件单独拎出来删掉或修复。要不然训练到后期突然读坏一张图报错整个epoch白跑。另外统一把所有图片转成RGB三通道避免混入灰度图导致ResNet输入维度对不上。2.2 数据增强策略与实现数据增强对农业图像分类不是可选项而是必需品。农田照片背景复杂同一种害虫在不同光照、不同叶片、不同拍摄角度下外观差异很大。不做增强模型很容易被某些颜色背景带偏换一批图片验证准确率就往下掉。我给的训练集transforms是这样的from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomVerticalFlip(p0.2), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.05), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里每个参数都有讲究。ResNet的输入尺寸固定为224x224验证阶段先Resize到256再做CenterCrop是标准做法能保证目标基本落在画面中央同时保留一点上下文信息。RandomResizedCrop的scale范围设成(0.8, 1.0)而不是默认的(0.08, 1.0)是因为害虫在照片里的占比通常不大如果裁剪过多很容易把目标切掉上限定为1.0用来偶尔做一次全图输入让模型不要只依赖局部特征。ColorJitter的hue只给0.05因为虫害诊断中颜色是重要的辨别信号色相抖动太大会把真实颜色信息冲掉。特别注意训练集和验证集千万不能用同一个transform。验证集一定不要做随机翻转和随机裁剪不然验证指标不稳定选最优模型都会受影响。太多新手在这里栽跟头训练集套个增强顺手复制到验证集上最后模型看起来训练loss很低一上线就露馅。2.3 自定义Dataset与DataLoader配置如果目录结构已经是ImageFolder格式那读取代码非常简洁from torchvision import datasets train_dataset datasets.ImageFolder(rootdataset/train, transformtrain_transform) val_dataset datasets.ImageFolder(rootdataset/val, transformval_transform)ImageFolder返回(图片张量, 类别索引)对自动按子文件夹名的字符串顺序分配索引。前面说的类别映射表在这里就有了用处最后输出结果时一定要回头核实对应关系。如果数据不是干净的目录结构比如标签存在CSV里每行是图片路径和标签字符串那就自己写Dataset类import pandas as pd from PIL import Image from torch.utils.data import Dataset class RicePestDataset(Dataset): def __init__(self, df, transformNone, class_to_idxNone): self.df df.reset_index(dropTrue) self.transform transform if class_to_idx is not None: self.class_to_idx class_to_idx else: self.class_to_idx {c: i for i, c in enumerate(sorted(df[label].unique()))} self.labels [self.class_to_idx[l] for l in self.df[label]] def __len__(self): return len(self.df) def __getitem__(self, idx): path self.df.loc[idx, path] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, self.labels[idx]DataLoader的参数我通常这样设置from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue, ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue, )batch_size根据显存调整。drop_lastTrue用于训练集能避免最后一个不完整batch在BatchNorm层造成统计偏差这个设置对loss稳定性有些帮助。num_workers在Windows上偶尔会引发多进程启动问题遇到的话把num_workers降到0或者把训练代码放进ifname main保护块。pin_memoryTrue在GPU训练时能降低数据传输等待实测有效。3. 模型构建与训练全流程3.1 加载预训练模型与替换分类头torchvision的加载方式在新版本中推荐传入weights参数我的写法是这样import torch import torch.nn as nn from torchvision import models def create_model(num_classes14): weights models.ResNet50_Weights.IMAGENET1K_V1 model models.resnet50(weightsweights) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return model这段代码的关键改动就在最后那一行。ResNet50的fc原始输出是1000现在改成14。in_features从这个层的输入特征数动态获取以后换ResNet18特征数是512也不用改代码。冻结主干的操作很简单遍历参数后把requires_grad设成Falsefor name, param in model.named_parameters(): if name.startswith(fc): param.requires_grad True else: param.requires_grad False这种策略适合数据集非常小的情况。但针对14类水稻虫害我最后采用的是两步走先冻结主干只训练fc层几个epoch让随机初始化的分类头先适应预训练特征然后解冻全部参数用一个很小的学习率微调整个网络。这样既避免一开始就全量微调破坏预训练好的权重又能让backbone后期适应农业图像的数据分布。3.2 损失函数、优化器与学习率调度分类任务首选交叉熵损失PyTorch里就一行criterion nn.CrossEntropyLoss()不需要手动做softmaxCrossEntropyLoss内部会自动完成LogSoftmax和NLLLoss的计算直接输入logits就行。优化器我推荐SGD加momentum配余弦退火调度器optimizer torch.optim.SGD( model.parameters(), lr1e-3, momentum0.9, weight_decay5e-4, ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, eta_min1e-6, )这里有个常见误区新初始化的fc层与预训练backbone不该用同一个学习率。随机初始化的fc层需要大步快跑预训练权重需要小步微调。如果共用同一个优化器可以通过分组参数解决optimizer torch.optim.SGD([ {params: model.fc.parameters(), lr: 1e-3}, {params: [p for n, p in model.named_parameters() if not n.startswith(fc)], lr: 1e-4}, ], momentum0.9, weight_decay5e-4)这么写的原因很直白两个不同初始化状态的参数用同一学习率更新要么让随机层收敛太慢要么让预训练层被冲得太厉害。分组设置让fc层快跑backbone稳走整体收敛路径平滑很多。SGD配合余弦退火在训练末段会自然降温帮助模型在损失曲面的低谷处落脚。如果你想节省调参时间用AdamW也可以lr设置在1e-4到3e-4之间weight_decay设0.01到0.05。但我的经验是SGD在小数据集上的最终精度往往比Adam系更好AdamW的优势是前期收敛速度快怎么选看你的时间和调参耐心。3.3 训练循环的核心代码训练循环是整个流程的重头戏。我习惯把单个epoch封装成函数def train_one_epoch(model, loader, optimizer, criterion, device): model.train() running_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device) labels labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc几个细节展开说一下。optimizer.zero_grad()每次迭代都必须调用否则梯度会累加惯用法是在backward之前清空上一轮梯度。torch.max(outputs, 1)返回两个值一个是最大值一个是最大值对应的索引索引就是预测类别。loss.item()取出的是一个标量这里乘上batch size再除以总样本数是为了计算整个epoch的平均loss。验证函数几乎一样区别是加model.eval()并且把推理代码包在torch.no_grad()里torch.no_grad() def validate(model, loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 for images, labels in loader: images images.to(device) labels labels.to(device) outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_accmodel.eval()会切换BatchNorm层和Dropout层的行为如果不调验证时的BatchNorm使用的统计量就不对结果和耍流氓没区别。torch.no_grad()能省掉梯度计算占用的显存推理阶段显存占用一般能减少将近一半。训练入口循环自然包含调度器步进和最优模型保存best_acc 0.0 for epoch in range(30): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_acc validate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch1:02d} | train_loss: {train_loss:.4f} | train_acc: {train_acc:.4f} | val_loss: {val_loss:.4f} | val_acc: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save({ epoch: epoch, model_state_dict: model.state_dict(), class_to_idx: train_dataset.class_to_idx, }, rice_pest_resnet50_best.pth)scheduler.step()要放在验证之后让下一轮优化器使用更新后的学习率。保存state_dict时连同类别映射一起保存这样后续推理时不需要自己重新对应索引和类别名。3.4 迁移学习微调时机与策略很多人的第一次尝试是加载预训练模型后直接全量训练不冻结任何层。这样也能收敛但有两个隐患一是数据量不足时预训练权重很快被破坏模型退化成重新随机初始化的效果二是loss曲线不稳定验证准确率上蹿下跳。更稳妥的做法是分两阶段。第一阶段冻结backbone只训练fc层。这个阶段一般10到15个epoch就够学习率可以设大一些因为随机初始化的fc层需要相对大步长去快速拟合特征分布。观察训练loss和验证loss都在持续下降说明分类头已经把特征映射掌握得差不多。第二阶段解冻全部参数把学习率降一个量级。整个网络一起微调去适应农业图像特有的统计规律。这一阶段往往不用跑太久10到20个epoch就能看到验证准确率出现明显提升同时训练loss还在缓慢下降。真正训练过头的时候你会发现验证准确率开始掉点这时候就该停了。这个两阶段策略背后的逻辑很容易理解。预训练模型的低层卷积已经学到大量通用视觉特征比如边缘、角点、纹理这些特征对水稻害虫图片依然适用不需要被新数据重新冲乱而网络高层的语义特征与数据分布关系密切需要重新调整。先固定低层、调整高层再逐步放开低层整个优化路径更加平滑比一开始就全量训练更容易收敛到好的极值点。4. 模型评估与问题排查实录4.1 评估指标与混淆矩阵分析训练结束不能只看验证集总准确率。水稻虫害数据很容易类别不均衡比如褐飞虱的数量可能是稻纵卷叶螟的好几倍。假设一个大类占总样本的60%一个只会预测那一类的模型准确率已经到了60%看着还行但其他虫害类别几乎没有识别能力。所以我习惯用整体准确率加上分类报告、混淆矩阵一起评估。分类报告可以用sklearn快速生成from sklearn.metrics import classification_report, confusion_matrix preds_all, labels_all [], [] for images, labels in test_loader: images images.to(device) outputs model(images) _, predicted torch.max(outputs, 1) preds_all.extend(predicted.cpu().numpy()) labels_all.extend(labels.numpy()) print(classification_report(labels_all, preds_all, target_namesclass_names)) cm confusion_matrix(labels_all, preds_all)分类报告里的precision、recall、f1-score能清晰反映每个类别的真实识别能力。对水稻虫害场景我更关心recall因为漏报的代价比误报高一种虫害没识别出来农户就错过了防治窗口宁可多报也不该漏。如果发现某个类别recall明显偏低说明它的图像特征容易被其他类别混淆这时要么回到数据增强想办法要么补充该类别的真实样本。混淆矩阵是找出哪两类彼此混淆的直接工具。假设稻纵卷叶螟经常被识别成稻螟蛉大多是幼虫体型和颜色相似也可能是这两类样本在训练集中的背景高度相似。出现这种情况先查看混淆的样本图片确认是标注错误还是特征实在太接近。如果是后者考虑补充不同虫龄、不同拍摄角度的样本或者在数据增强中加入随机遮挡逼着模型更关注主体特征。4.2 过拟合与欠拟合的识别处理训练曲线是判断模型状态的第一手信息。过拟合的典型表现是训练集loss一路走低验证集loss在某个epoch后开始回升或震荡最终验证准确率停滞。处理手段按优先级排增加数据增强强度比如把RandomResizedCrop的scale范围放宽到(0.6, 1.0)加入RandomErasing模拟遮挡调大weight_decay从5e-4试到1e-3甚至更大减少全量微调的epoch数提前早停检查类别样本数量对样本过少的类别做重采样考虑换小模型比如用ResNet34替代ResNet50小模型在数据不足时反而优势更明显。欠拟合的表现刚好相反训练集loss和验证集loss都居高不下验证准确率一直低位徘徊。常见原因是训练不充分、学习率设置不合理或者模型容量不够。优先检查学习率先往上一档调一档试试另外确认一下是不是误把所有主干参数都设成requires_gradFalse导致只有fc层在学习。还要警惕一个容易忽略的点输入预处理做错了比如忘记做Normalize。模型输入数值范围不对loss会漂在一个异常高的值附近怎么调优化器都没用很多人在这上面浪费了大半天时间。4.3 常见报错与解决办法速查整个训练过程中踩过的坑整理成速查表按报错现象、常见原因、解决办法三列对齐报错现象常见原因解决办法加载预训练模型时报size mismatch分类头类别数与预训练权重1000类不一致先执行model.fc nn.Linear(2048, 14)再load_state_dict或者直接使用torchvision自带预训练接口RuntimeError: CUDA out of memorybatch_size过大输入尺寸超过预期降低batch_size确认输入已经Resize到256并CenterCrop到224考虑pin_memory和梯度累积DataLoader多进程启动报错Windows下缺少ifname main保护num_workers0训练代码放进main函数或者把num_workers置为0PIL.UnidentifiedImageError原数据混入损坏图片文件写脚本批量Image.verify()扫描异常文件单独剔除准确率一直低loss不降输入忘记归一化或类别标签映射错位检查transform是否包含Normalize打印一批图片张量和标签确认对应关系训练loss正常验证loss偏高过拟合或者验证集错误地使用了随机增强增加数据增强强度检查val_transform是否包含随机操作这个表看起来简单每一条背后都是实际项目中的血泪教训。尤其是第一条size mismatch很多人看到报错直接懵了其实原理上就是模型结构的最后一层和checkpoint里的权重形状不匹配先把分类头改掉再加载权重就解决了。4.4 抽样验证与推理阶段准备训练好模型只是第一步真正能交付还要把模型放到真实图片上跑通。我通常最后写一个小脚本从测试集外抽几张真实图片做快速推理from PIL import Image def predict_image(img_path, model, transform, class_names, device): model.eval() img Image.open(img_path).convert(RGB) img_tensor transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs model(img_tensor) _, pred torch.max(outputs, 1) return class_names[pred.item()], torch.softmax(outputs, dim1).max().item()这里把预测类别和置信度一起输出是有意为之。如果模型对某张图的置信度低于0.6就不直接输出害虫种类而是标记为待人工确认。这个置信度阈值在农业落地项目中很有用它会大幅减少误报对农户的误导。如果后续还有部署需求可以用TorchScript把模型导出再接入推理服务或边缘设备。5. 项目复盘与扩展建议5.1 数据质量优先于模型结构这个项目走到最后我最大的体会是模型结构不是瓶颈数据质量才是。14类水稻虫害的分类难度并不算高但田间照片的采集规范对最终效果影响极大。有些类别准确率看着很高细看原因是训练集里全是同一品种、同一角度、同一光照的照片模型学到的是场景记忆而不是通用特征。我后来用同一套代码处理另一组野外采集数据室内拍摄的测试集上去到98%准确率换到室外实拍照片直接掉到80%出头差距全在训练集与测试集的拍摄条件不同。所以从一开始就应该把数据按拍摄环境、光照条件、季节分层划分不要在同一个文件夹里混入大量高度相似的连拍图。最好给每个类别都留出一个独立的野外验证集防止模型被训成只认某一块田的模型。5.2 后续功能扩展方向一个完整的水稻虫害识别系统远不止训练一个14类分类器。模型训完后还有很多可以做的事把模型导出成ONNX或TorchScript部署到后端服务提供一个图片上传接口用推理框架做加速单张图片推理耗时尽量控制在几十毫秒内才能满足实地使用对低置信度样本做主动学习定期收集真实用户反馈中的难例补充标注后做增量训练结合水稻生长周期和虫害发生规律让系统给出更符合农学实际的预警结果。这些扩展已经超出了单纯的模型训练但如果你从一开始就把数据读取、模型构建、推理函数拆开写后面接任何服务框架都不需要大改这也是我在这个项目中最推荐的工程习惯。最后分享一个我自己养成的小习惯每次训练前固定随机种子并且在代码里把train/val/test的图片数量统计表打印出来。这样隔几天回来复现结果时能清楚知道自己用的是哪个版本的数据模型效果波动到底来自模型改动还是数据变动。跑通整个流程不难难的是把流程每个环节都变成可回溯、可复现、可持续改进的闭环这才是真正能落地到农业生产环境里的深度学习项目。