小数据集害虫图像分类实战：从PyTorch迁移学习到YOLOv5

简介：这份图像分类数据包面向需要庄稼害虫识别训练数据的开发者、学生或算法工程师，包含蛀虫、健康无虫、螨虫等4个类别，提供训练集620张与验证集53张。数据按文件夹保存为ImageFolder格式，可直接用PyTorch的ImageFolder加载，无需额外预处理；同样适配YOLOv5分类训练，适合快速验证农业害虫识别模型。资源共676个文件，主体为673张JPG图片，另含1个JSON分类字典、1个可视化预览py脚本和1张展示png，压缩包仅53.89MB，目录按train/test划分，下载和部署都很快捷。附带的py脚本无需修改即可运行，随机传入4张图片生成预览并保存到当前目录，能直观核对图像、标签和分类字典是否一致，避免训练前踩坑。该资源已有276人浏览学习，对入门图像分类或农业病虫害识别、做课程设计或论文实验的读者来说，是一份开箱即用的标准数据集，能有效节省数据采集、清洗与整理时间。

1. 四类庄稼害虫图像分类：小数据集如何做出可用基线

农业场景里，虫害识别的第一难题不是模型选哪个，而是找不到带标签的数据。实际项目里经常只有几十到几百张现场照片，分类类别也远没有公开数据集那么规整。这个包提供的四分类害虫数据，包含620张训练图和53张验证图，类别覆盖蛀虫、健康无虫、螨虫等典型状态，目录已经按train/test分好，用ImageFolder读取后不需要额外写转换脚本。对于刚接触图像分类算法的人，可以拿它跑通迁移学习全流程；对于有经验的工程师，又可以用它验证yolov5训练自己的数据集时的目录规范。数据量虽然不大，但恰好够说明“小样本分类”里最常见的坑：过拟合、类别分布不均和验证集过小。

2. 数据组织方式与ImageFolder加载原理：目录就是标签

2.1 解压后的目录结构与类别文件夹的含义

打开数据包后，核心目录是data，下面两个子目录分别是train和test。很多从标注平台导出的数据会写成train/val/test三份，但这份数据把验证集直接命名为test。在图像分类任务里，验证集和测试集的本质区别在于是否参与模型选择；用test目录做验证时，要注意最终报告的精度不能代表真实泛化精度，因为它可能被你重复用来调参。实操时可以先当验证集用，报告结果时说明这一点。

目录展开后大致是：

data/ ├── train/ │ ├── borer/ │ ├── healthy/ │ ├── mite/ │ └── other/ └── test/ ├── borer/ ├── healthy/ ├── mite/ └── other/

类名我这里用英文示意。原包里的文件夹名可能保留中文或拼音，这不影响加载，但影响后续模型输出和脚本可读性，建议训练前统一转成小写英文。每类文件夹下是一组JPG图片，文件名里能看到类似xxx.rf.xxx.jpg的中间标记，这是从Roboflow平台导出时常见的命名风格，说明图片已经过平台侧的处理或增强。文件后缀大小写混用，读取时要注意。

2.2 用命令行快速核对训练集与验证集数量

拿到数据先别急着写训练脚本，用两个命令确认图片数量和类别数量：

find data/train -type f -iname '*.jpg' | wc -l find data/test -type f -iname '*.jpg' | wc -l

第一个命令统计train目录下所有.jpg文件，-iname是不区分大小写的文件名匹配，能同时抓住.JPG和.jpg。如果统计结果和标注一致，训练集620、验证集53，说明文件没有缺漏；如果少于这个数，就要检查是否混入了非图片格式或隐藏文件。类别数量可以通过ls data/train看到四个文件夹，四分类是确定的。53张验证图、每类平均约13张，这个规模做精度评估会很抖，后面训练时建议同时记录每次验证的类别准确率，而不是只看总体准确率。

2.3 ImageFolder 的类别索引是怎么生成的

PyTorch 的ImageFolder是读取这类目录型数据最直接的工具。它不读标注文件，而是按文件夹名字自动生成类别映射。核心用法：

from torchvision.datasets import ImageFolder train_ds = ImageFolder("data/train") print(train_ds.classes) print(train_ds.class_to_idx)

classes是文件夹名称列表，顺序按字典序排列，而不是按文件在磁盘上的原始顺序。例如borer, healthy, mite, other会被映射为 0、1、2、3。这个映射顺序决定了模型输出 logits 的通道对应哪个类。如果训练脚本里硬编码类别 ID，换机器或换数据集时很容易错位；正确做法是每次从train_ds.class_to_idx里读取并保存到 JSON，后续推理时也读取同一个 JSON。

ImageFolder构造函数里有几个常用参数：root是包含分类子目录的根路径；transform接收图片变换；target_transform可以修改标签。取图片时可以用train_ds.samples[i]，返回(图片路径, 标签)二元组；直接下标train_ds[i]返回(tensor, label)，但如果没指定transform，也可能是 PIL 图像。小数据集排查问题时，建议先不带transform遍历一遍，确认所有图片都能正常打开，避免读图报错淹没后面的训练问题。

2.4 JSON 分类字典文件：与 ImageFolder 保持一致

原包里附带了一个 JSON 文件，内容应该就是四分类的类别字典。为了训练时统一，可以按下面方式生成一份新的 JSON：

import json import os train_dir = "data/train" classes = sorted([d for d in os.listdir(train_dir) if os.path.isdir(os.path.join(train_dir, d))]) class_to_idx = {cls: idx for idx, cls in enumerate(classes)} with open("classes.json", "w", encoding="utf-8") as f: json.dump(class_to_idx, f, indent=2, ensure_ascii=False) print(class_to_idx)

注意这里sorted要与ImageFolder的排序逻辑一致。如果 JSON 里的映射和class_to_idx不一致，推理阶段很容易出现类别错位且没有任何报错。生成后的文件内容类似：

这个 JSON 不仅给 PyTorch 训练用，后面转 YOLOv5 分类模式时，同样需要它来核对类别顺序。训练前跑一次sorted(os.listdir())和ImageFolder.class_to_idx对比，能提前发现文件名里包含奇怪字符导致排序不一致的坑。

2.5 训练集和验证集比例偏少时的应对

训练集620、验证集53，验证占比不足10%，比常用80/20划分更紧。直接用这个验证集做模型选择，精度波动可能很大。常见做法是先按现有 train/test 跑一版，确认流程没问题，再合并两边数据做一次分层采样，得到新的训练集和验证集。分层采样的关键是按类别分别 shuffle 后抽取，避免某类全部跑进验证集，这里给一个通用函数：

import random from collections import defaultdict from torch.utils.data import Subset def stratified_split(dataset, val_ratio=0.2, seed=42): idx_by_class = defaultdict(list) for i, (_, label) in enumerate(dataset): idx_by_class[label].append(i) train_idx, val_idx = [], [] rng = random.Random(seed) for label, indices in idx_by_class.items(): rng.shuffle(indices) n_val = max(1, int(len(indices) * val_ratio)) val_idx.extend(indices[:n_val]) train_idx.extend(indices[n_val:]) return Subset(dataset, train_idx), Subset(dataset, val_idx)

dataset可以传入一个合并后的ConcatDataset，也可以直接传入ImageFolder的实例。这里的Subset不复制图片，只是记录索引，训练时依然按原路径读图，内存占用很低。重新划分后，每个类在验证集中大约保留20%个样本，620张图会得到约124张验证图，比原来的53张可靠不少。要注意的是，新划分得到的结果和原包自带的 test 标签不再一致，做对照实验时不要混用。

3. 基于预训练模型搭建四分类器：迁移学习训练与验证

3.1 为什么小样本四分类仍然适合用深度模型

620张图对深度学习来说不多，但图像分类算法在 ImageNet 上预训练出的模型已经掌握了大量纹理、边缘和局部形状特征。害虫与健康作物的差异，很多体现在体表纹理和区域颜色分布上，这些特征与 ImageNet 中的自然图像有重叠，所以用迁移学习比从零训练更稳。常见的选择是 ResNet 或 EfficientNet，这里以 ResNet18 为例：它只有约1100万参数，在一张中等显卡上几分钟就能跑完一个 epoch，适合快速验证数据质量。

如果机器更紧张，也可以换mobilenet_v3_small或efficientnet_b0，替换分类输出维度即可。对四分类任务来说，模型容量反而不是瓶颈，瓶颈往往是数据规模太小导致的过拟合，因此后续会用较轻量的数据增强和早停来控制。

3.2 数据加载与增强：Resize、翻转与归一化

下面代码是训练阶段推荐的一组预处理：

from torchvision import transforms from torchvision.datasets import ImageFolder from torch.utils.data import DataLoader load_size = 224 train_tf = transforms.Compose([ transforms.Resize((load_size, load_size)), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(10), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) val_tf = transforms.Compose([ transforms.Resize((load_size, load_size)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) train_ds = ImageFolder("data/train", transform=train_tf) val_ds = ImageFolder("data/test", transform=val_tf) train_loader = DataLoader(train_ds, batch_size=16, shuffle=True, num_workers=4, drop_last=True) valid_loader = DataLoader(val_ds, batch_size=16, shuffle=False, num_workers=4)

Resize((224, 224))的作用是把所有图片统一为正方形输入，避免不同相机尺寸导致 batch 拼接失败。RandomHorizontalFlip和RandomRotation属于轻量增强，能提高模型对拍摄角度变化的抗性，但旋转角度不要太大，农作物害虫本身有方向性，旋转90度以上可能产生无意义的样本。Normalize里使用的是 ImageNet 数据统计值，迁移学习加载预训练权重时原则上要继续沿用，不能随意替换成自己的均值和标准差，否则会破坏权重分布。drop_last=True在验证集上不要用，因为会直接丢失最后一批；训练集上则没问题，还可以避免最后一个 batch 过小引起 BatchNorm 统计量抖动。

3.3 替换全连接层与两阶段优化

ResNet18 最后的全连接层输入是512维，输出是1000类。把输出改成4类：

import torch.nn as nn from torchvision import models model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) in_features = model.fc.in_features model.fc = nn.Linear(in_features, 4) print(model.fc)

model.fc.in_features读取的是原分类层的输入维度，这样即使换成其他系列模型，代码也不用改。weights=models.ResNet18_Weights.IMAGENET1K_V1表示加载官方 ImageNet 预训练权重，比pretrained=True这种旧写法更显式，避免未来版本移除旧接口时报警告。

训练时先固定主干，只更新分类头：

for param in model.parameters(): param.requires_grad = False for param in model.fc.parameters(): param.requires_grad = True criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

这阶段相当于把预训练特征当作固定向量的分类器，几十个 epoch 内就能快速收敛。然后解冻全部层，用更小的学习率微调：

for param in model.parameters(): param.requires_grad = True optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

两阶段优化是迁移学习里很常用的做法。直接上来全部微调，在大学习率下容易把预训练特征“洗掉”，导致小数据集上反而比只训分类头更差。解冻后学习率降到 1e-4，相当于在原始特征附近做局部搜索。如果显存不够或时间紧，可以只解冻最后两个残差块，其余保持冻结，这时需要按模块名选择参数，代码会稍复杂，但思路相同。

3.4 训练循环、验证精度与最佳权重保存

一个完整的最小训练循环如下：

import torch from tqdm import tqdm device = "cuda" if torch.cuda.is_available() else "cpu" model = model.to(device) for epoch in range(15): model.train() total_loss = 0.0 for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}"): images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) # 验证 model.eval() correct = 0 total = 0 with torch.no_grad(): for images, labels in valid_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, preds = torch.max(outputs, dim=1) correct += (preds == labels).sum().item() total += labels.size(0) acc = correct / total print(f"train_loss={total_loss / len(train_ds):.4f} " f"val_acc={acc:.4f}") torch.save(model.state_dict(), f"pest_{epoch+1:02d}.pt")

torch.max(outputs, dim=1)返回每个样本的最大 logit 位置，这就是预测类别。上面的训练循环是在第一阶段冻结主干下跑的，requires_grad=False时 BatchNorm 的 running_mean/var 也会保持冻结，即使model.train()也不会被当前 batch 更新。验证时一定要切到model.eval()，否则 BatchNorm 统计量会被当前 batch 污染，精度会出现明显波动。

15 个 epoch 对第一阶段偏多，实际跑的时候可以在第5个 epoch 就观察验证集是否不再上升。对小数据集来说，验证集只有53张，一个样本的误判就会让准确率跳约1.9%，所以看到精度小幅波动不用急着调参，先把学习率和 epoch 固定住，多跑几个随机种子看稳定水平。

4. 面向YOLOv5分类模式的数据集适配与训练

4.1 YOLOv5分类训练对目录命名的要求

用 yolo 系列做图像分类，经典仓库是ultralytics/yolov5的classify目录。YOLOv5 分类数据集的默认要求是train和val两个同级的类文件夹根目录，文件夹内再按类别分子目录。当前数据包给的是train和test，需要做一个小映射。最简单的方式是创建软链接，避免复制多一份数据：

cd data ln -s test val

执行后data目录下同时出现train和val，YOLOv5 的classify/train.py会直接读取它们。如果操作系统不支持软链接，也可以用cp -r test val，但白占一份磁盘空间。对于620张图片来说无所谓，但对于后续更大的数据集建议用ln -s。

注意，不要直接改test文件夹的名字，因为 PyTorch 训练脚本里还要用data/test做验证评估，两边同时需要。软链接可以两全其美。

4.2 使用 classify/train.py 启动训练

克隆并进入yolov5仓库后，直接运行：

python classify/train.py --model yolov5s-cls.pt \ --data ../data \ --epochs 30 \ --img 224 \ --batch 16

--model指定分类预训练模型，yolov5s-cls.pt是在 ImageNet 上训练过的分类权重；--data指向包含train/val的根目录；--epochs对于620张训练图可以设30～50；--img是输入分辨率，一般224足够，较大分辨率能保留更多害虫细节但会明显增加训练时间；--batch在显存足够时设16或32，数值过小会导致 BatchNorm 统计不稳定。

启动后日志里会出现每个 epoch 的验证准确率。训练结束后权重保存在runs/train-cls/exp/weights/best.pt，这份权重可以直接用于后续预测。如果只想验证数据格式有没有问题，可以加--epochs 1，一个 epoch 内就能跑完，避免浪费时间在错误配置上。

4.3 验证与评估：用 val.py 和混淆矩阵检查类别表现

训练完成后，可以用官方验证脚本跑一遍，得到整体精度和混淆矩阵：

python classify/val.py --weights runs/train-cls/exp/weights/best.pt \ --data ../data

脚本会输出 Top-1 Accuracy、Top-5 Accuracy 以及每个类别的详细指标。之所以要看混淆矩阵，是因为害虫数据里不同类之间可能存在视觉相似，例如某类螨虫和蛀虫在图像上可能都表现为叶片斑点，整体准确率0.85时，某两类的互相混淆可能已经很高。混淆矩阵图保存在runs/val-cls/exp/下，文件名通常是confusion_matrix.png。

如果发现某个类完全预测不对，优先检查该类在训练集和验证集中的样本数分布。四分类中如果某一类训练图明显少于其他类，需要先做类间复制或图像增强，而不是急着加网络深度。另外验证集只有53张，混淆矩阵的单个格子里放的是百分比还是绝对数量要看清，避免把个别样本的误差解读成系统性错误。

4.4 YOLOv8分类模式的兼容性说明

YOLOv8 的官方包也支持分类任务，数据目录格式和 YOLOv5 一样，训练入口更简单：

yolo classify train model=yolov8n-cls.pt data=../data epochs=30 imgsz=224 batch=16

yolov8n-cls.pt是轻量分类预训练权重，训练完成后同样在runs/classify/train下输出best.pt。如果你在用数据包练习 yolov8 训练自己的数据集，这个目录结构可以直接复用。需要注意 YOLOv8 会对data参数做路径解析，建议使用绝对路径或完整相对路径，避免在子目录里运行时找不到数据集。

这里补充一点，YOLOv5 和 YOLOv8 的分类验证指标都依赖val目录，软链接方式在两者之间都能生效。后面如果要把data/test同时保留，就不要在源码里写死val路径，尽量用命令行参数切换。

5. 快速验证与可视化：随机抽样脚本与预测结果检查

5.1 一条命令可视化验证集样本

原包附带的可视化脚本思路是随机从数据集中取4张图，拼成一张网格后保存。我自己常用的写法是：

import random import matplotlib.pyplot as plt from torchvision.datasets import ImageFolder from torchvision.utils import make_grid ds = ImageFolder("data/test") samples = random.sample(range(len(ds)), 4) images = [ds[i][0] for i in samples] grid = make_grid(images, nrow=2) plt.figure(figsize=(6, 6)) plt.imshow(grid.permute(1, 2, 0)) plt.axis("off") plt.savefig("random_check.png", bbox_inches="tight")

ds[i][0]返回的是未规范化 PIL 图像，make_grid内部会做边界填充，nrow=2表示一行两列，4张图正好两行。这个脚本的重点是快速确认验证集图片是否损坏、类别目录是否对应预期。如果make_grid报类型错误，多半是ds传入时已经带了transform导致得到 Tensor，这时需要把ImageFolder的transform参数留空，或在网格生成前手动还原归一化。

5.2 从置信度分布里找误判样本

模型训练完，不要只看准确率。把验证集所有样本预测一遍，按置信度升序排列，通常排在最前面的就是最容易误判的样本。

import torch from torchvision import transforms from torchvision.datasets import ImageFolder val_tf = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) ds = ImageFolder("data/test", transform=val_tf) model.eval() results = [] with torch.no_grad(): for i, (img, label) in enumerate(ds): logits = model(img.unsqueeze(0)) prob = torch.softmax(logits, dim=1) conf, pred = prob.max(dim=1) results.append((conf.item(), ds.classes[pred.item()], ds.classes[label], ds.samples[i][0])) results.sort(key=lambda x: x[0]) for conf, pred, truth, path in results[:5]: print(f"{conf:.2f} pred={pred:<10} truth={truth:<10} path={path}")

这段代码对验证集所有样本做一次前向推理，softmax后取最大置信度和对应类别。对置信度最低的5个样本，可以打开原图看是标注错误、图像模糊，还是两类本身视觉边界模糊。如果发现大量低置信度样本集中在某一类，就要考虑给该类增加更多训练样本或调整分类阈值。53张验证集虽然小，但足够做一次系统性的误判复盘。

本文还有配套的精品资源，点击获取