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 + torchvision

PyTorch在迁移学习场景中的优势太明显了。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类水稻虫害分类任务，我习惯按照下面这条主线推进：

1. 确认原始数据形态，整理成train/val/test三类目录，每个类别单独一个子目录；
2. 用ImageFolder或自定义Dataset构建训练集和验证集的读取逻辑；
3. 在transforms里给训练集加入随机裁剪、翻转、颜色抖动等增强，验证集只做固定尺寸和归一化；
4. 加载预训练ResNet，替换最后全连接层，决定是否冻结主干；
5. 配置损失函数、优化器、学习率调度器，把训练循环封装成函数；
6. 每个epoch结束后在验证集上评估，保存验证集表现最好的权重；
7. 训练完成后加载最优权重，在测试集上输出分类报告和混淆矩阵等指标。

这个流程把“一定会用到的环节”和“需要按实际情况调整的环节”分开了。数据增强策略可以根据类别分布与图片质量调整，冻结主干与否可以根据预训练模型与数据的相似程度调整，但整体框架是稳定的。下面就是把这条路线一步一步落地的实际操作。

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(p=0.5), transforms.RandomVerticalFlip(p=0.2), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.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(root='dataset/train', transform=train_transform) val_dataset = datasets.ImageFolder(root='dataset/val', transform=val_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, transform=None, class_to_idx=None): self.df = df.reset_index(drop=True) 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_size=32, shuffle=True, num_workers=4, pin_memory=True, drop_last=True, ) val_loader = DataLoader( val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True, )

batch_size根据显存调整。drop_last=True用于训练集，能避免最后一个不完整batch在BatchNorm层造成统计偏差，这个设置对loss稳定性有些帮助。num_workers在Windows上偶尔会引发多进程启动问题，遇到的话把num_workers降到0，或者把训练代码放进ifname== 'main'保护块。pin_memory=True在GPU训练时能降低数据传输等待，实测有效。

3. 模型构建与训练全流程

3.1 加载预训练模型与替换分类头

torchvision的加载方式在新版本中推荐传入weights参数，我的写法是这样：

import torch import torch.nn as nn from torchvision import models def create_model(num_classes=14): weights = models.ResNet50_Weights.IMAGENET1K_V1 model = models.resnet50(weights=weights) 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设成False：

for 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()

不需要手动做softmax，CrossEntropyLoss内部会自动完成LogSoftmax和NLLLoss的计算，直接输入logits就行。

优化器我推荐SGD加momentum，配余弦退火调度器：

optimizer = torch.optim.SGD( model.parameters(), lr=1e-3, momentum=0.9, weight_decay=5e-4, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=30, eta_min=1e-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}, ], momentum=0.9, weight_decay=5e-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_acc

model.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(f"Epoch {epoch+1: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_names=class_names)) cm = confusion_matrix(labels_all, preds_all)

分类报告里的precision、recall、f1-score能清晰反映每个类别的真实识别能力。对水稻虫害场景，我更关心recall，因为漏报的代价比误报高，一种虫害没识别出来，农户就错过了防治窗口，宁可多报也不该漏。如果发现某个类别recall明显偏低，说明它的图像特征容易被其他类别混淆，这时要么回到数据增强想办法，要么补充该类别的真实样本。

混淆矩阵是找出哪两类彼此混淆的直接工具。假设稻纵卷叶螟经常被识别成稻螟蛉，大多是幼虫体型和颜色相似，也可能是这两类样本在训练集中的背景高度相似。出现这种情况先查看混淆的样本图片，确认是标注错误还是特征实在太接近。如果是后者，考虑补充不同虫龄、不同拍摄角度的样本，或者在数据增强中加入随机遮挡，逼着模型更关注主体特征。

4.2 过拟合与欠拟合的识别处理

训练曲线是判断模型状态的第一手信息。

过拟合的典型表现是训练集loss一路走低，验证集loss在某个epoch后开始回升或震荡，最终验证准确率停滞。处理手段按优先级排：

1. 增加数据增强强度，比如把RandomResizedCrop的scale范围放宽到(0.6, 1.0)，加入RandomErasing模拟遮挡；
2. 调大weight_decay，从5e-4试到1e-3甚至更大；
3. 减少全量微调的epoch数，提前早停；
4. 检查类别样本数量，对样本过少的类别做重采样；
5. 考虑换小模型，比如用ResNet34替代ResNet50，小模型在数据不足时反而优势更明显。

欠拟合的表现刚好相反，训练集loss和验证集loss都居高不下，验证准确率一直低位徘徊。常见原因是训练不充分、学习率设置不合理，或者模型容量不够。优先检查学习率，先往上一档调一档试试，另外确认一下是不是误把所有主干参数都设成requires_grad=False，导致只有fc层在学习。

还要警惕一个容易忽略的点：输入预处理做错了，比如忘记做Normalize。模型输入数值范围不对，loss会漂在一个异常高的值附近，怎么调优化器都没用，很多人在这上面浪费了大半天时间。

4.3 常见报错与解决办法速查

整个训练过程中踩过的坑，整理成速查表，按报错现象、常见原因、解决办法三列对齐：

这个表看起来简单，每一条背后都是实际项目中的血泪教训。尤其是第一条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, dim=1).max().item()

这里把预测类别和置信度一起输出是有意为之。如果模型对某张图的置信度低于0.6，就不直接输出害虫种类，而是标记为待人工确认。这个置信度阈值在农业落地项目中很有用，它会大幅减少误报对农户的误导。如果后续还有部署需求，可以用TorchScript把模型导出，再接入推理服务或边缘设备。

5. 项目复盘与扩展建议

5.1 数据质量优先于模型结构

这个项目走到最后，我最大的体会是模型结构不是瓶颈，数据质量才是。14类水稻虫害的分类难度并不算高，但田间照片的采集规范对最终效果影响极大。有些类别准确率看着很高，细看原因是训练集里全是同一品种、同一角度、同一光照的照片，模型学到的是场景记忆而不是通用特征。我后来用同一套代码处理另一组野外采集数据，室内拍摄的测试集上去到98%准确率，换到室外实拍照片直接掉到80%出头，差距全在训练集与测试集的拍摄条件不同。

所以从一开始，就应该把数据按拍摄环境、光照条件、季节分层划分，不要在同一个文件夹里混入大量高度相似的连拍图。最好给每个类别都留出一个独立的野外验证集，防止模型被训成只认某一块田的模型。

5.2 后续功能扩展方向

一个完整的水稻虫害识别系统远不止训练一个14类分类器。模型训完后还有很多可以做的事：

1. 把模型导出成ONNX或TorchScript，部署到后端服务，提供一个图片上传接口；
2. 用推理框架做加速，单张图片推理耗时尽量控制在几十毫秒内，才能满足实地使用；
3. 对低置信度样本做主动学习，定期收集真实用户反馈中的难例，补充标注后做增量训练；
4. 结合水稻生长周期和虫害发生规律，让系统给出更符合农学实际的预警结果。

这些扩展已经超出了单纯的模型训练，但如果你从一开始就把数据读取、模型构建、推理函数拆开写，后面接任何服务框架都不需要大改，这也是我在这个项目中最推荐的工程习惯。

最后分享一个我自己养成的小习惯：每次训练前固定随机种子，并且在代码里把train/val/test的图片数量统计表打印出来。这样隔几天回来复现结果时，能清楚知道自己用的是哪个版本的数据，模型效果波动到底来自模型改动还是数据变动。跑通整个流程不难，难的是把流程每个环节都变成可回溯、可复现、可持续改进的闭环，这才是真正能落地到农业生产环境里的深度学习项目。