简介小麦种子图像分类数据集面向计算机视觉学习者、研究者以及农业智能化项目开发者提供约2000张已标注图片覆盖4个品种类别数据已经过预处理可直接送入分类网络训练与验证避免从零处理原始图像的重复劳动。资源共2000个文件包含1998张jpg图片、1个Python可视化脚本及1个json标签配置文件整体压缩包约225.64MB7z格式目录结构清晰便于按需取用。包内训练集、测试集已分开存放图片也按同一类别归入对应文件夹运行show脚本可快速预览和检查样本分布json文件则明确标注类别索引省去自行整理标签的环节。目前已有221人学习下载适合用于分类网络练手、模型对比试验也可作为数据预处理与可视化流程的参考基准资源还附有作者关于图像分类/分割网络改进及计算机视觉完整项目的博文入口方便进一步扩展实践。1. 小麦种子图像分类数据集从 2,000 张已标注图片到分类网络输入这个小麦种子图像分类数据集的规模不大但很完整约 2,000 张已标注图片、4 个品种类别已经按训练集和测试集分好目录并预先做了统一的预处理。文件名本身带有类别前缀例如 aca315_0219.jpg、buckguapo_0283.jpg配合 json 里的类别映射基本可以做到开箱即用。对做图像分类的团队来说它的价值一是省去采集标注的脏活二是可以作为自建数据集的目录规范和标注模板对准备跑通图像分类算法 baseline 的新手也能在半天内完成从读图到训练再到验证的闭环。2. 数据集结构与 JSON 标注训练/测试划分与类别映射拿到压缩包后先别急着跑训练第一步是把目录结构和标注文件摸清楚。这个数据集的关键特征是「已标注、已划分、已预处理」这三个状态决定了你后续代码怎么写也决定了哪些环节可以省、哪些环节仍然不能省。2.1 目录布局按类别分目录省掉手写标签逻辑我一般会先观察顶层目录常见整理方式是按 train/test 分开各自下面再按类别建子目录。参考给出的文件名整理后的结构大致是dataset/ ├── train/ │ ├── aca315/ │ │ ├── aca315_0219.jpg │ │ └── aca315_0222.jpg │ ├── buckguapo/ │ │ ├── buckguapo_0282.jpg │ │ ├── buckguapo_0283.jpg │ │ └── buckguapo_0284.jpg │ └── buckmeteoro/ │ ├── buckmeteoro_0081.jpg │ ├── buckmeteoro_0095.jpg │ ├── buckmeteoro_0206.jpg │ └── buckmeteoro_0215.jpg ├── test/ │ ├── aca315/ │ │ └── ... │ ├── buckguapo/ │ │ └── ... │ └── buckmeteoro/ │ └── ... └── labels.json这是根据文件名做的常规推断实际目录名可能有细微差别但思路一样。按类名建目录的好处是 torchvision 的 ImageFolder 可以直接使用不需要额外维护 image_path 到 label_index 的字典。文件名里的前缀和目录名保持一致所以即使 json 中只给了文件清单也能用 os.path 反查标签。这个设计在复现时非常省事自己整理数据时也值得沿用。labels.json 里最重要的信息是类别顺序。常见字段如下表但不同发布者的字段名不一定相同所以第一步永远是打印 keys() 确认。常见字段示例说明class_names[aca315, buckguapo, buckmeteoro, ...]类别顺序必须与 ImageFolder 的目录字母序一致train[train/aca315/aca315_0219.jpg, ...]训练集文件清单test[test/buckguapo/buckguapo_0283.jpg, ...]测试集文件清单类别顺序很容易被忽略但直接影响结果。ImageFolder 会按子目录名的字母序给类别编号如果 json 里的 class_names 顺序与目录字母序不一致轻则标签错位重则整个模型学不到东西。我的做法是先把 json 加载出来再用目录统计做交叉验证。import json from collections import Counter from pathlib import Path with open(labels.json, r, encodingutf-8) as f: meta json.load(f) # 第一步先看顶层有哪些键再决定后续解析逻辑 print(meta.keys()) # 用目录结构做交叉验证不依赖 json 字段名 train_dir Path(train) counter Counter(p.parent.name for p in train_dir.rglob(*.jpg)) print(counter)这段代码里meta.keys()帮你确认 json 结构是文件清单模式还是简单映射模式Path.rglob(*.jpg)递归匹配 train 目录下所有 jpg 图片p.parent.name取图片所在子目录名也就是类别名。Counter 统计每个类别的图片数量一是用来确认类别个数是否与描述一致二是为后续判断类别是否平衡做准备。如果 counter 的键数量与 class_names 长度不一致说明划分文件里漏了某个类别先解决数据问题再谈训练。2.2 show 脚本可视化样本人工确认标签没有错位原资源带了 show 脚本常见运行方式是python show_data.py --data_root dataset --num_show 8脚本逻辑通常是从每个类别取若干样本用 matplotlib 画成网格标题显示类别名。这里给一个等价实现import matplotlib.pyplot as plt from PIL import Image def show_samples(paths, labels, class_names, num8): rows num // 4 fig, axes plt.subplots(rows, 4, figsize(12, 3 * rows)) for i, ax in enumerate(axes.flat): img Image.open(paths[i]).convert(RGB) ax.imshow(img) ax.set_title(class_names[labels[i]]) ax.axis(off) plt.tight_layout() plt.show()这里用 PIL 打开图片后先转 RGB防止灰度图在 matplotlib 里变成二维数组导致显示异常class_names[labels[i]]确保标题显示的是类别名而不是整数索引。运行这一步能发现两类问题一是文件名和图像内容是否匹配二是同一类目录下是否存在损坏图、灰度图、比例异常图。如果种子个体在图片里占比很小建议手工裁切局部放大检查因为后续 Resize 到 224x224 后细小纹路可能在缩放中被抹掉。2.3 训练/测试划分与验证集缺口这个数据集已经划分好训练集和测试集这是优点但也带来了一个隐患没有独立的验证集。按照严格流程验证集用于调参和选模型测试集只允许最终评估一次。如果直接拿 test 反复试 learning rate、调整增强参数最后报告的指标会偏乐观因为模型已经隐式见过测试集的信息。常见做法是从 train 里再切出一部分作为验证集比如 10% 到 20%。2000 张图总量不算多切 15% 左右比较合理训练集仍剩 1400 张左右足够做迁移学习。切分时保持每个类别的比例避免某个类在验证集中一个样本都没有。这一步可以用 sklearn 的 train_test_split 完成也可以直接用目录遍历后按文件名移动。在项目初期把这个验证集固定下来后续所有实验都只在这个验证集上比较最后再用 test 做一次干净评估。3. 预处理与 DataLoader 参数让 2,000 张图片训练得更稳数据已经经过预处理可以直接作为分类网络输入。但「已预处理」不等于「不需要 transform」。通常预处理只是统一了尺寸、格式和后缀到了训练代码里仍然要做 ToTensor 和 Normalize。这一章讲清楚 transform 怎么写、DataLoader 参数怎么设以及小数据集上最常踩的坑。3.1 先检查图片的实际尺寸和通道不要因为「预处理过」就跳过检查。用几行 Python 快速验证from PIL import Image img Image.open(train/aca315/aca315_0219.jpg) print(img.size, img.mode)如果输出是(224, 224)和RGB那么预处理已经统一了分辨率和通道后续 transform 可以直接沿用。如果发现有的图是 224x224有的是 512x384那么 transform 里的 Resize 仍然是必须的不能因为「预处理过」三个字就省略掉。mode如果是L或RGBA要在 transform 之前统一转成 RGB否则后面 Normalize 的三个通道统计量会直接报错或者产生错误数值。我一般会在自定义 Dataset 的__getitem__里加一句.convert(RGB)成本极低但能省掉一批偶发的通道问题。3.2 Transform 参数小物体分类的增强要克制对小麦种子这类小物体图像分类数据集增强策略和普通场景分类不同。种子在图像中占比小、背景简单过强的 RandomResizedCrop 很容易裁到背景导致模型学到背景纹理而不是种子轮廓。我一般用弱增强from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])训练集和验证集的变换必须分开写验证集只做 Resize、ToTensor、Normalize不做随机翻转和旋转否则每次评估结果都会因为随机性而抖动。Resize 固定到 224x224 是为了匹配预训练模型的输入尺寸RandomHorizontalFlip 和 RandomRotation 是对种子外观影响较小的增强旋转 10 度以内不会改变品种特征。数据总量只有 2000 张增强过强反而会让模型在小样本下更难收敛。Normalize 使用 ImageNet 统计量因为后面要加载 ImageNet 预训练权重输入分布必须与预训练一致。3.3 ImageFolder 与 DataLoader 参数由于目录按类别分好直接用 ImageFolder 加载是最省事的做法from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_ds ImageFolder(rootdataset/train, transformtrain_transform) val_ds ImageFolder(rootdataset/test, transformval_transform) train_loader DataLoader( train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue, ) val_loader DataLoader( val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue, )ImageFolder 会按子目录名的字母序为类别分配从 0 开始的标签这个顺序必须和 labels.json 中的 class_names 一致建议在训练前打印train_ds.class_to_idx做一次确认。shuffleTrue是为了避免每个 batch 内全是同一类别drop_lastTrue是丢弃最后一个不足 batch 的样本在 BatchNorm 层较多时能减少统计量抖动。参数层面的参考值如下。参数建议值为什么batch_size16 或 322000 张样本32 一次迭代约 62 个 batch足够 BatchNorm 统计显存小就用 16shuffletrain True / test False防止同类样本扎堆进入同一 batchnum_workers4~8提高图像解码吞吐Windows 上报 BrokenPipeError 就改回 0 或 2pin_memory训练时 TrueCUDA 拷贝更高效CPU 训练时无效drop_last训练时 True避免最后一个小 batch 导致统计量抖动num_workers不是越大越好它取决于 CPU 核数和磁盘读取速度。2000 张图的数据量并不大4 个 worker 通常已经够用如果训练时 CPU 被打满而 GPU 利用率不高不要急着加 worker先看是不是 transform 里的 Resize 太耗时。4. 从 ResNet 到 Transformer 图像分类训练脚本与超参数配置4 分类、2000 张样本、已划分训练测试集这个规模适合用「预训练权重 微调」策略。直接用随机初始化训练一个 ResNet18 也能跑但需要更多 epoch 和更强的数据增强而使用 ImageNet 预训练权重后15 个 epoch 左右就能稳定收敛。4.1 模型选型2000 张数据不要一上来就上大模型经常看到有人拿 2000 张图直接训 ViT 或 Swin Transformer结果在没有大规模预训练时不仅收敛慢还比 ResNet 差。这里给出选型经验表模型结构特点2000 张数据上的建议ResNet18层数少、不易过拟合首选 baselineResNet50容量更大细节特征更多搭配弱增强并增大 weight_decayViT-S/16全局注意力需要大量数据只建议加载预训练权重后微调最后几层对于小麦种子这类纹理差异体现在整体形状和表面纹路的图像ResNet 系列的卷积归纳偏置在小样本上更占便宜。如果想尝试 Transformer 图像分类模型建议把它当作第二阶段的实验而不是第一个跑通的模型。另外要注意图像的旋转方向如果种子本身有明显的方向性RandomRotation 就需要关闭这一点可以在 show 脚本可视化时直接判断。4.2 训练脚本迁移学习、损失函数与优化器加载预训练 ResNet18 并替换最后一个全连接层import torch import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, 4) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW( model.parameters(), lr1e-4, weight_decay1e-4 ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max15, eta_min1e-6 )替换全连接层后模型输出维度变成 4对应 labels.json 中的 4 个类别。CrossEntropyLoss 在 PyTorch 中已经包含 softmax所以最后一层不需要额外加 LogSoftmax。AdamW 的 lr 从 1e-4 起步比较稳如果是从头训练lr 可以提高到 3e-4但过拟合风险会显著增加。CosineAnnealingLR 的 T_max 设为总 epoch 数让学习率在一个周期内从 1e-4 平滑降到 1e-6适合小数据集微调。训练主循环可以这样写model.cuda() model.train() for epoch in range(15): total_loss, correct 0.0, 0 for x, y in train_loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() total_loss loss.item() * x.size(0) correct (logits.argmax(1) y).sum().item() scheduler.step() avg_loss total_loss / len(train_ds) acc correct / len(train_ds) print(fepoch {epoch1}: loss{avg_loss:.4f}, acc{acc:.4f})每个 epoch 输出平均 loss 和训练准确率能及时发现 loss 不下降或 acc 抖动。loss.item() * x.size(0)是为了用样本数加权避免最后一个 batch 大小不一致导致平均 loss 有偏。logits.argmax(1)取每个样本预测类别。scheduler.step()放在每个 epoch 结束后如果放在 optimizer.step() 之后学习率衰减就会快一个 batch属于常见低级错误。4.3 类别不平衡与采样器2000 张图按 4 个类别划分如果某些品种图像只有 300 张某些有 700 张直接用 CrossEntropyLoss 会偏向样本多的一类。可以先统计类别分布差值超过 2 倍就要处理。最简单的办法是给损失函数加类别权重class_counts torch.tensor([...], dtypetorch.float) # 从统计结果手动填入 weights 1.0 / class_counts criterion nn.CrossEntropyLoss(weightweights.cuda())也可以使用 WeightedRandomSampler 在取样本时平衡类别from torch.utils.data import WeightedRandomSampler labels torch.tensor(train_ds.targets) class_counts torch.bincount(labels) sample_weights 1.0 / class_counts[labels] sampler WeightedRandomSampler( sample_weights, num_sampleslen(sample_weights), replacementTrue ) train_loader DataLoader( train_ds, batch_size32, samplersampler, num_workers4 )weight 方式更简单逻辑上等价于对少数类的 loss 进行放大sampler 方式让少数类样本在每次采样中出现概率更高。两者选一个即可不要同时用。使用 sampler 后DataLoader 不能再设置shuffleTrue否则会报参数冲突有 sampler 时 DataLoader 会在每个 epoch 内部重新打乱顺序取满指定数量所以也不需要再手动 shuffle。5. 用混淆矩阵和类别 F1 验证品种分类效果训练结束后如果只看测试准确率往往会漏掉两个问题小类被大类吞掉、某两个品种系统性混淆。种子图像相邻品种通常颜色或形状接近checkpoint 保存后先做一次完整的推理再输出混淆矩阵和分类报告。5.1 生成分类报告并定位易混类别这段代码可以直接复用from sklearn.metrics import classification_report, confusion_matrix all_preds, all_labels [], [] model.eval() with torch.no_grad(): for x, y in val_loader: x x.cuda() logits model(x) all_preds logits.argmax(dim1).cpu().tolist() all_labels y.cpu().tolist() print(classification_report(all_labels, all_preds, target_namestrain_ds.classes)) print(confusion_matrix(all_labels, all_preds))分类报告里的 macro avg 比 accuracy 更能反映小类表现混淆矩阵对角线是正确分类数非对角线是易混对。如果某两个类别反复互混第一反应不是改网络结构而是回到 show 脚本看这两类的原始图像是否存在光照不一致、背景色接近、拍摄角度差异等问题。常见做法是挑出 5 张错分样本把原图、真实标签、预测标签拼在一起确认是标注边界问题还是特征本身不可分。如果要进一步量化多分类全局能力可以计算每个类别的 AUC。对 4 分类做 One-vs-Rest 展开用 sklearn 的 label_binarize 配合 roc_auc_score 就可以完成。保存模型时建议用验证集 loss 而不是最后一个 epoch每轮结束比较 val_loss只保留历史最佳权重避免最后几轮已经过拟合时把模型覆盖掉。本文还有配套的精品资源点击获取
简介小麦种子图像分类数据集面向计算机视觉学习者、研究者以及农业智能化项目开发者提供约2000张已标注图片覆盖4个品种类别数据已经过预处理可直接送入分类网络训练与验证避免从零处理原始图像的重复劳动。资源共2000个文件包含1998张jpg图片、1个Python可视化脚本及1个json标签配置文件整体压缩包约225.64MB7z格式目录结构清晰便于按需取用。包内训练集、测试集已分开存放图片也按同一类别归入对应文件夹运行show脚本可快速预览和检查样本分布json文件则明确标注类别索引省去自行整理标签的环节。目前已有221人学习下载适合用于分类网络练手、模型对比试验也可作为数据预处理与可视化流程的参考基准资源还附有作者关于图像分类/分割网络改进及计算机视觉完整项目的博文入口方便进一步扩展实践。1. 小麦种子图像分类数据集从 2,000 张已标注图片到分类网络输入这个小麦种子图像分类数据集的规模不大但很完整约 2,000 张已标注图片、4 个品种类别已经按训练集和测试集分好目录并预先做了统一的预处理。文件名本身带有类别前缀例如 aca315_0219.jpg、buckguapo_0283.jpg配合 json 里的类别映射基本可以做到开箱即用。对做图像分类的团队来说它的价值一是省去采集标注的脏活二是可以作为自建数据集的目录规范和标注模板对准备跑通图像分类算法 baseline 的新手也能在半天内完成从读图到训练再到验证的闭环。2. 数据集结构与 JSON 标注训练/测试划分与类别映射拿到压缩包后先别急着跑训练第一步是把目录结构和标注文件摸清楚。这个数据集的关键特征是「已标注、已划分、已预处理」这三个状态决定了你后续代码怎么写也决定了哪些环节可以省、哪些环节仍然不能省。2.1 目录布局按类别分目录省掉手写标签逻辑我一般会先观察顶层目录常见整理方式是按 train/test 分开各自下面再按类别建子目录。参考给出的文件名整理后的结构大致是dataset/ ├── train/ │ ├── aca315/ │ │ ├── aca315_0219.jpg │ │ └── aca315_0222.jpg │ ├── buckguapo/ │ │ ├── buckguapo_0282.jpg │ │ ├── buckguapo_0283.jpg │ │ └── buckguapo_0284.jpg │ └── buckmeteoro/ │ ├── buckmeteoro_0081.jpg │ ├── buckmeteoro_0095.jpg │ ├── buckmeteoro_0206.jpg │ └── buckmeteoro_0215.jpg ├── test/ │ ├── aca315/ │ │ └── ... │ ├── buckguapo/ │ │ └── ... │ └── buckmeteoro/ │ └── ... └── labels.json这是根据文件名做的常规推断实际目录名可能有细微差别但思路一样。按类名建目录的好处是 torchvision 的 ImageFolder 可以直接使用不需要额外维护 image_path 到 label_index 的字典。文件名里的前缀和目录名保持一致所以即使 json 中只给了文件清单也能用 os.path 反查标签。这个设计在复现时非常省事自己整理数据时也值得沿用。labels.json 里最重要的信息是类别顺序。常见字段如下表但不同发布者的字段名不一定相同所以第一步永远是打印 keys() 确认。常见字段示例说明class_names[aca315, buckguapo, buckmeteoro, ...]类别顺序必须与 ImageFolder 的目录字母序一致train[train/aca315/aca315_0219.jpg, ...]训练集文件清单test[test/buckguapo/buckguapo_0283.jpg, ...]测试集文件清单类别顺序很容易被忽略但直接影响结果。ImageFolder 会按子目录名的字母序给类别编号如果 json 里的 class_names 顺序与目录字母序不一致轻则标签错位重则整个模型学不到东西。我的做法是先把 json 加载出来再用目录统计做交叉验证。import json from collections import Counter from pathlib import Path with open(labels.json, r, encodingutf-8) as f: meta json.load(f) # 第一步先看顶层有哪些键再决定后续解析逻辑 print(meta.keys()) # 用目录结构做交叉验证不依赖 json 字段名 train_dir Path(train) counter Counter(p.parent.name for p in train_dir.rglob(*.jpg)) print(counter)这段代码里meta.keys()帮你确认 json 结构是文件清单模式还是简单映射模式Path.rglob(*.jpg)递归匹配 train 目录下所有 jpg 图片p.parent.name取图片所在子目录名也就是类别名。Counter 统计每个类别的图片数量一是用来确认类别个数是否与描述一致二是为后续判断类别是否平衡做准备。如果 counter 的键数量与 class_names 长度不一致说明划分文件里漏了某个类别先解决数据问题再谈训练。2.2 show 脚本可视化样本人工确认标签没有错位原资源带了 show 脚本常见运行方式是python show_data.py --data_root dataset --num_show 8脚本逻辑通常是从每个类别取若干样本用 matplotlib 画成网格标题显示类别名。这里给一个等价实现import matplotlib.pyplot as plt from PIL import Image def show_samples(paths, labels, class_names, num8): rows num // 4 fig, axes plt.subplots(rows, 4, figsize(12, 3 * rows)) for i, ax in enumerate(axes.flat): img Image.open(paths[i]).convert(RGB) ax.imshow(img) ax.set_title(class_names[labels[i]]) ax.axis(off) plt.tight_layout() plt.show()这里用 PIL 打开图片后先转 RGB防止灰度图在 matplotlib 里变成二维数组导致显示异常class_names[labels[i]]确保标题显示的是类别名而不是整数索引。运行这一步能发现两类问题一是文件名和图像内容是否匹配二是同一类目录下是否存在损坏图、灰度图、比例异常图。如果种子个体在图片里占比很小建议手工裁切局部放大检查因为后续 Resize 到 224x224 后细小纹路可能在缩放中被抹掉。2.3 训练/测试划分与验证集缺口这个数据集已经划分好训练集和测试集这是优点但也带来了一个隐患没有独立的验证集。按照严格流程验证集用于调参和选模型测试集只允许最终评估一次。如果直接拿 test 反复试 learning rate、调整增强参数最后报告的指标会偏乐观因为模型已经隐式见过测试集的信息。常见做法是从 train 里再切出一部分作为验证集比如 10% 到 20%。2000 张图总量不算多切 15% 左右比较合理训练集仍剩 1400 张左右足够做迁移学习。切分时保持每个类别的比例避免某个类在验证集中一个样本都没有。这一步可以用 sklearn 的 train_test_split 完成也可以直接用目录遍历后按文件名移动。在项目初期把这个验证集固定下来后续所有实验都只在这个验证集上比较最后再用 test 做一次干净评估。3. 预处理与 DataLoader 参数让 2,000 张图片训练得更稳数据已经经过预处理可以直接作为分类网络输入。但「已预处理」不等于「不需要 transform」。通常预处理只是统一了尺寸、格式和后缀到了训练代码里仍然要做 ToTensor 和 Normalize。这一章讲清楚 transform 怎么写、DataLoader 参数怎么设以及小数据集上最常踩的坑。3.1 先检查图片的实际尺寸和通道不要因为「预处理过」就跳过检查。用几行 Python 快速验证from PIL import Image img Image.open(train/aca315/aca315_0219.jpg) print(img.size, img.mode)如果输出是(224, 224)和RGB那么预处理已经统一了分辨率和通道后续 transform 可以直接沿用。如果发现有的图是 224x224有的是 512x384那么 transform 里的 Resize 仍然是必须的不能因为「预处理过」三个字就省略掉。mode如果是L或RGBA要在 transform 之前统一转成 RGB否则后面 Normalize 的三个通道统计量会直接报错或者产生错误数值。我一般会在自定义 Dataset 的__getitem__里加一句.convert(RGB)成本极低但能省掉一批偶发的通道问题。3.2 Transform 参数小物体分类的增强要克制对小麦种子这类小物体图像分类数据集增强策略和普通场景分类不同。种子在图像中占比小、背景简单过强的 RandomResizedCrop 很容易裁到背景导致模型学到背景纹理而不是种子轮廓。我一般用弱增强from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees10), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])训练集和验证集的变换必须分开写验证集只做 Resize、ToTensor、Normalize不做随机翻转和旋转否则每次评估结果都会因为随机性而抖动。Resize 固定到 224x224 是为了匹配预训练模型的输入尺寸RandomHorizontalFlip 和 RandomRotation 是对种子外观影响较小的增强旋转 10 度以内不会改变品种特征。数据总量只有 2000 张增强过强反而会让模型在小样本下更难收敛。Normalize 使用 ImageNet 统计量因为后面要加载 ImageNet 预训练权重输入分布必须与预训练一致。3.3 ImageFolder 与 DataLoader 参数由于目录按类别分好直接用 ImageFolder 加载是最省事的做法from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_ds ImageFolder(rootdataset/train, transformtrain_transform) val_ds ImageFolder(rootdataset/test, transformval_transform) train_loader DataLoader( train_ds, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue, ) val_loader DataLoader( val_ds, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue, )ImageFolder 会按子目录名的字母序为类别分配从 0 开始的标签这个顺序必须和 labels.json 中的 class_names 一致建议在训练前打印train_ds.class_to_idx做一次确认。shuffleTrue是为了避免每个 batch 内全是同一类别drop_lastTrue是丢弃最后一个不足 batch 的样本在 BatchNorm 层较多时能减少统计量抖动。参数层面的参考值如下。参数建议值为什么batch_size16 或 322000 张样本32 一次迭代约 62 个 batch足够 BatchNorm 统计显存小就用 16shuffletrain True / test False防止同类样本扎堆进入同一 batchnum_workers4~8提高图像解码吞吐Windows 上报 BrokenPipeError 就改回 0 或 2pin_memory训练时 TrueCUDA 拷贝更高效CPU 训练时无效drop_last训练时 True避免最后一个小 batch 导致统计量抖动num_workers不是越大越好它取决于 CPU 核数和磁盘读取速度。2000 张图的数据量并不大4 个 worker 通常已经够用如果训练时 CPU 被打满而 GPU 利用率不高不要急着加 worker先看是不是 transform 里的 Resize 太耗时。4. 从 ResNet 到 Transformer 图像分类训练脚本与超参数配置4 分类、2000 张样本、已划分训练测试集这个规模适合用「预训练权重 微调」策略。直接用随机初始化训练一个 ResNet18 也能跑但需要更多 epoch 和更强的数据增强而使用 ImageNet 预训练权重后15 个 epoch 左右就能稳定收敛。4.1 模型选型2000 张数据不要一上来就上大模型经常看到有人拿 2000 张图直接训 ViT 或 Swin Transformer结果在没有大规模预训练时不仅收敛慢还比 ResNet 差。这里给出选型经验表模型结构特点2000 张数据上的建议ResNet18层数少、不易过拟合首选 baselineResNet50容量更大细节特征更多搭配弱增强并增大 weight_decayViT-S/16全局注意力需要大量数据只建议加载预训练权重后微调最后几层对于小麦种子这类纹理差异体现在整体形状和表面纹路的图像ResNet 系列的卷积归纳偏置在小样本上更占便宜。如果想尝试 Transformer 图像分类模型建议把它当作第二阶段的实验而不是第一个跑通的模型。另外要注意图像的旋转方向如果种子本身有明显的方向性RandomRotation 就需要关闭这一点可以在 show 脚本可视化时直接判断。4.2 训练脚本迁移学习、损失函数与优化器加载预训练 ResNet18 并替换最后一个全连接层import torch import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, 4) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW( model.parameters(), lr1e-4, weight_decay1e-4 ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max15, eta_min1e-6 )替换全连接层后模型输出维度变成 4对应 labels.json 中的 4 个类别。CrossEntropyLoss 在 PyTorch 中已经包含 softmax所以最后一层不需要额外加 LogSoftmax。AdamW 的 lr 从 1e-4 起步比较稳如果是从头训练lr 可以提高到 3e-4但过拟合风险会显著增加。CosineAnnealingLR 的 T_max 设为总 epoch 数让学习率在一个周期内从 1e-4 平滑降到 1e-6适合小数据集微调。训练主循环可以这样写model.cuda() model.train() for epoch in range(15): total_loss, correct 0.0, 0 for x, y in train_loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() total_loss loss.item() * x.size(0) correct (logits.argmax(1) y).sum().item() scheduler.step() avg_loss total_loss / len(train_ds) acc correct / len(train_ds) print(fepoch {epoch1}: loss{avg_loss:.4f}, acc{acc:.4f})每个 epoch 输出平均 loss 和训练准确率能及时发现 loss 不下降或 acc 抖动。loss.item() * x.size(0)是为了用样本数加权避免最后一个 batch 大小不一致导致平均 loss 有偏。logits.argmax(1)取每个样本预测类别。scheduler.step()放在每个 epoch 结束后如果放在 optimizer.step() 之后学习率衰减就会快一个 batch属于常见低级错误。4.3 类别不平衡与采样器2000 张图按 4 个类别划分如果某些品种图像只有 300 张某些有 700 张直接用 CrossEntropyLoss 会偏向样本多的一类。可以先统计类别分布差值超过 2 倍就要处理。最简单的办法是给损失函数加类别权重class_counts torch.tensor([...], dtypetorch.float) # 从统计结果手动填入 weights 1.0 / class_counts criterion nn.CrossEntropyLoss(weightweights.cuda())也可以使用 WeightedRandomSampler 在取样本时平衡类别from torch.utils.data import WeightedRandomSampler labels torch.tensor(train_ds.targets) class_counts torch.bincount(labels) sample_weights 1.0 / class_counts[labels] sampler WeightedRandomSampler( sample_weights, num_sampleslen(sample_weights), replacementTrue ) train_loader DataLoader( train_ds, batch_size32, samplersampler, num_workers4 )weight 方式更简单逻辑上等价于对少数类的 loss 进行放大sampler 方式让少数类样本在每次采样中出现概率更高。两者选一个即可不要同时用。使用 sampler 后DataLoader 不能再设置shuffleTrue否则会报参数冲突有 sampler 时 DataLoader 会在每个 epoch 内部重新打乱顺序取满指定数量所以也不需要再手动 shuffle。5. 用混淆矩阵和类别 F1 验证品种分类效果训练结束后如果只看测试准确率往往会漏掉两个问题小类被大类吞掉、某两个品种系统性混淆。种子图像相邻品种通常颜色或形状接近checkpoint 保存后先做一次完整的推理再输出混淆矩阵和分类报告。5.1 生成分类报告并定位易混类别这段代码可以直接复用from sklearn.metrics import classification_report, confusion_matrix all_preds, all_labels [], [] model.eval() with torch.no_grad(): for x, y in val_loader: x x.cuda() logits model(x) all_preds logits.argmax(dim1).cpu().tolist() all_labels y.cpu().tolist() print(classification_report(all_labels, all_preds, target_namestrain_ds.classes)) print(confusion_matrix(all_labels, all_preds))分类报告里的 macro avg 比 accuracy 更能反映小类表现混淆矩阵对角线是正确分类数非对角线是易混对。如果某两个类别反复互混第一反应不是改网络结构而是回到 show 脚本看这两类的原始图像是否存在光照不一致、背景色接近、拍摄角度差异等问题。常见做法是挑出 5 张错分样本把原图、真实标签、预测标签拼在一起确认是标注边界问题还是特征本身不可分。如果要进一步量化多分类全局能力可以计算每个类别的 AUC。对 4 分类做 One-vs-Rest 展开用 sklearn 的 label_binarize 配合 roc_auc_score 就可以完成。保存模型时建议用验证集 loss 而不是最后一个 epoch每轮结束比较 val_loss只保留历史最佳权重避免最后几轮已经过拟合时把模型覆盖掉。本文还有配套的精品资源点击获取
