简介小麦种子图像分类数据集是一份面向图像分类任务学习与研究的已标注数据资源适合深度学习初学者或农业表型分析人员用于训练和验证分类网络。数据集对原始图片完成了预处理共划分4个类别训练集与测试集已按类目分好可直接作为分类模型输入省去数据清洗和文件整理环节。包内另附json标签配置文件和Python可视化脚本用户可运行脚本查看各类别图像与标签对应情况便于快速理解数据分布。整个压缩包约225.64MB共2000个文件其中1998张jpg图像、1个py脚本和1个json配置文件结构简洁清晰。该资源已有223人学习下载对需要现成带标注作物图像数据、快速开展图像分类实验的读者来说具有较好的实用价值。1. 2000张已标注的小麦种子图像够不够做出能用的分类模型很多工程师一听到“只有2000张已标注的图像分类数据集”第一反应是先扩数据。小麦种子分类这个场景其实不用——类别少常见做法按品种分为五类上下种质差异主要靠纹理、色泽和轮廓不像自动驾驶那样依赖海量场景。用 ImageNet 预训练权重做迁移学习2000张已标注数据足以在一两小时内刷出可用的 baseline继续调优后完全能到九成上下。真正决定这件事能不能落地的不是样本数而是标签干不干净、验证集划分严谨不严谨、预处理和后端部署有没有对齐。这篇笔记按我自己做小样本图像分类的顺序讲从数据体检到训练、参数再到几个悄悄吃掉准确率的坑给手里正好有小样本种质图像、想快速拿到可信结果的工程师和农业科研人员参考。这里不玩玄学全部是可复现的步骤。2. 小麦种子数据集的标签与分布训练前最该做的三件事2.1 用脚本统计类别、尺寸与损坏文件标签和数据分布是图像分类项目里最容易被跳过的环节。很多人的做法是解压完数据集直接开训等到训练完才发现少数类只有几十张图或者某个目录里混入了损坏文件白白浪费几小时。我一般会先花十分钟做数据体检。常见目录结构是按类别建子文件夹比如train/红硬麦/xxx.jpg也有单层文件夹配一张labels.csv的。第一步先写几行脚本把整体情况打出来from pathlib import Path from collections import Counter from PIL import Image data_dir Path(./wheat_seed/images) # 改成你的数据目录 exts {.jpg, .jpeg, .png, .bmp, .tif} files [p for p in data_dir.rglob(*) if p.suffix.lower() in exts] print(样本总数:, len(files)) labels [p.parent.name if p.parent ! data_dir else unknown for p in files] counter Counter(labels) for k, v in counter.items(): print(f{k}: {v} ({v / len(files):.1%})) bad [] for f in files: try: with Image.open(f) as im: im.load() except Exception: bad.append(f) print(损坏文件数:, len(bad))这段代码做了三件事统计文件数、按父目录名统计类别占比、逐张验证图像能否真正解码。注意exts里必须带.tif很多种质图像是显微镜或扫描仪导出的 TIFF不带后缀匹配会把它们漏掉。损坏检测里的im.load()会把像素数据真正读入内存遇到截断的 JPEG 或 0 字节文件会直接抛异常比只看扩展名可靠得多。如果数据集是单层目录配 CSV那就用pandas读标签表import pandas as pd df pd.read_csv(./wheat_seed/labels.csv) print(df.head()) print(df[label].value_counts(normalizeTrue))这里要提醒一句CSV 里的文件名和磁盘上的文件名必须 join 之后统一核对一遍Windows 和 Linux 的路径分隔符不一致的问题在跨平台解压数据集时经常出现。2.2 看占比类不平衡会让整体准确率虚高类别分布打出来后你第一眼要看的不是总样本数而是占比。如果某个品种占了六成甚至七成一个“全部猜它”的哑模型也能拿到六七十的整体准确率。训练起来你会发现 loss 降得很快但小品种始终被吞。在小数据集上我更推荐用类别权重而不是过采样。2000张已经不多过采样等于把少数类反复复制模型很容易把少数类背下来验证集上看着不错换一批新图又掉回去。类别权重的做法很简单在第 4 章训练循环里把CrossEntropyLoss加上weight参数就行如果某个类只有 80 张不要急着扩样本先试权重多数情况下能把整体准确率拉回正常区间。同时可以顺手看一下图像尺寸分布。种子图片如果是扫描仪扫出来的通常分辨率统一如果混入了不同来源的图尺寸会七零八落。模型输入固定 224x224太小的图比如低于 100 的上采样后会变成糊成一团的色块这类样本最好直接筛掉不要保留在训练集里。2.3 按样本组划分给相似图片留一道隔离带这一步是最容易被忽略、影响力却最大的划分陷阱。小麦种子数据集里同一粒麦子往往拍了多张照片或者同一批样本在同一光照条件下连续拍摄。如果直接把所有图片随机shuffle后按比例划分同源照片会一边一张地出现在训练集和验证集里验证集等于开卷考试准确率高得离谱一到现场换一批不相关的麦子立刻打回原形。解决方法是先用文件名或 CSV 里的sample_id提取分组键按组而不是按单张图划分。假设文件名长得像WHT-038_a.jpg、WHT-038_b.jpg那么WHT-038就是同源组import random from collections import defaultdict files_by_group defaultdict(list) for f in files: group f.stem.split(_)[0] # WHT-038_a - WHT-038 files_by_group[group].append(f) groups list(files_by_group.keys()) random.seed(42) random.shuffle(groups) split int(len(groups) * 0.8) train_groups set(groups[:split]) val_groups set(groups[split:]) train_files [f for g in train_groups for f in files_by_group[g]] val_files [f for g in val_groups for f in files_by_group[g]] print(训练组数:, len(train_groups), 样本数:, len(train_files)) print(验证组数:, len(val_groups), 样本数:, len(val_files))注意如果你的文件名没有明确的样本 ID就按拍摄批次的文件夹划分。原则是同一场景、同一批次、同一颗麦穗出镜的图像必须整组进同一边。stem.split(_)[0]的写法只适用于“下划线前是同源 ID”这类命名命名规则不同时改成对应的截取方式即可。关键点是random.seed(42)必须在 shuffle 之前保证别人复现你的划分结果是一模一样的。3. 用迁移学习跑通小麦种子图像分类ResNet50 与数据增强的组合3.1 先把数据装进 Dataset尺寸、通道与 transform 一起写死有了划分结果下一步就是把图片和标签组织成 PyTorch 的Dataset。种子图像基本都是三通道彩色图但偶尔混入灰度图或带透明通道的 PNG所以在__getitem__里统一convert(RGB)是最稳妥的。import torch from torch.utils.data import Dataset from PIL import Image class WheatSeedDataset(Dataset): def __init__(self, file_list, labels, transformNone): self.file_list file_list self.labels labels self.transform transform def __len__(self): return len(self.file_list) def __getitem__(self, idx): image Image.open(self.file_list[idx]).convert(RGB) label self.labels[idx] if self.transform: image self.transform(image) return image, labelDataset本身不做归一化归一化全部交给transform。这里把通道处理写在数据读取层比在训练循环里判断通道数要清爽得多。transform是预处理的关键也是训练和验证必须分开的地方from torchvision import transforms train_tf transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.02), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) 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]), ])先Resize(256)再RandomCrop(224)等于在中心裁剪基础上加了小幅平移扰动相当于一种免费数据增强。如果你直接用Resize(224)模型每次看到的都是完全相同的构图在小数据集上更容易过拟合。mean和std必须是 ImageNet 预训练权重对应的那组不要自己重算否则等于把预训练特征空间篡改了。3.2 用 ResNet50 预训练权重做迁移学习先冻结只训分类头关于选哪个网络我一般会用 ResNet50 而不是最新的大模型。2000张数据量不适合从零训练 ViT而纯用最新的图像分类模型权重往往是为 ImageNet 的 1000 类设计直接用不了必须换分类头。ResNet50 的预训练特征对纹理、边缘、光泽这类低层视觉模式非常稳健对种子表皮这种目标特别合适。import torch import torch.nn as nn from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) num_classes 5 # 改成数据集实际类别数 model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) fc_in model.fc.in_features model.fc nn.Sequential( nn.Dropout(p0.2), nn.Linear(fc_in, num_classes) ) # 先冻结骨干只训练分类头 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True model.to(device) # 打印确认只有 fc 部分参数需要梯度 for name, param in model.named_parameters(): if param.requires_grad: print(可训练:, name)weightsmodels.ResNet50_Weights.IMAGENET1K_V1是当前 torchvision 的推荐写法别再用旧版的pretrainedTrue那个参数在新版本里被移除了。分类头从单层Linear换成Dropout Linear对 2000 张的小数据集能明显抑制分类层的过拟合。第一遍训练只解冻分类头骨干全部冻结这样做有两个好处一是参数量骤减训练速度快二是预训练特征不会被小数据集带偏。3.3 旋转翻转够用就行RandomErasing 这类增强反而破坏种子纹理数据增强不是越多越好。种子成像时朝向本来就随机所以RandomRotation(15°)和RandomHorizontalFlip是性价比最高的增强光照和色温会改变表皮高光ColorJitter也很关键。但有两类增强要特别小心。RandomResizedCrop默认的scale(0.08, 1.0)会随机裁得很小种子主体可能被裁掉一半模型最后学到的全是背景纹理。RandomErasing会直接擦除一块区域对自动驾驶场景可能有用但对种子的腹沟、表皮纹路这种决定品种的细粒度特征擦一块就等于毁掉关键证据。增强的原则是“保持类别不变的扰动”种子类别的核心差异在纹理和轮廓过强的几何扭曲会破坏这些特征。如果哪天你要在验证集上复现训练效果也记得用val_tf而不是训练增强否则验证结果永远在抖。4. 训练参数怎么设学习率、早停与收敛判据4.1 训练参数的常规组合小数据集图像分类的参数组合不同人的习惯差别很大但对 2000 张已标注数据来说下面这组是我常用的起步值参数冻结骨干阶段解冻微调阶段说明优化器AdamWAdamWAdamW 配合 weight_decay 比 Adam 更稳初始学习率3e-4 ~ 5e-41e-4 ~ 3e-4骨干是预训练的学习率大了会毁掉特征batch size16 或 3216 或 32样本只有2000大 batch 更容易过拟合max_epochs30~5020~40冻结阶段基本 30 轮以内收敛early stopping patience8~108~12验证集小、波动大patience 太短会误停4.2 用学习率扫描找到起步值迁移学习最常见的学习率错误是直接把 ResNet 从零训练时的1e-3拿来用。冻结阶段骨干不动只有分类头是全新初始化的1e-3很容易让 loss 开场就飙高。我习惯先做一个极简的学习率扫描def lr_scan(model, loader, optimizer, base_lr1e-6, max_lr5e-3, steps20): model.train() loader_iter iter(loader) lrs, losses [], [] lr base_lr for i in range(steps): try: x, y next(loader_iter) except StopIteration: loader_iter iter(loader) x, y next(loader_iter) x, y x.to(device), y.to(device) lr base_lr * (max_lr / base_lr) ** (i / (steps - 1)) for g in optimizer.param_groups: g[lr] lr optimizer.zero_grad() loss nn.functional.cross_entropy(model(x), y) loss.backward() optimizer.step() lrs.append(lr) losses.append(loss.item()) return lrs, losses这个函数用同一个数据迭代器在1e-6到5e-3之间按对数间隔扫 20 个点每个点跑一步。返回后你观察 loss 曲线loss 开始明显下降的那个点再往左一格就是合适的初始学习率。种子分类任务里答案通常落在3e-4附近。扫描用的 batch 不需要大16 张就够目的是看梯度方向不是看最终精度。4.3 完整训练循环早停、最优权重保存有了学习率和数据加载就可以跑主循环。我喜欢把训练和评估写成两个独立函数避免验证阶段误开梯度import torch import torch.nn.functional as F def train_one_epoch(model, loader, optimizer, criterion): model.train() total_loss, correct, total 0.0, 0, 0 for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() out model(x) loss criterion(out, y) loss.backward() optimizer.step() total_loss loss.item() * x.size(0) correct (out.argmax(1) y).sum().item() total x.size(0) return total_loss / total, correct / total torch.no_grad() def evaluate(model, loader): model.eval() total_loss, correct, total 0.0, 0, 0 for x, y in loader: x, y x.to(device), y.to(device) out model(x) loss F.cross_entropy(out, y) total_loss loss.item() * x.size(0) correct (out.argmax(1) y).sum().item() total x.size(0) return total_loss / total, correct / total optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4) criterion torch.nn.CrossEntropyLoss() max_epochs 40 patience 10 best_acc 0.0 wait 0 for epoch in range(max_epochs): tr_loss, tr_acc train_one_epoch(model, train_loader, optimizer, criterion) val_loss, val_acc evaluate(model, val_loader) print(fepoch {epoch1:02d} train_loss{tr_loss:.4f} val_loss{val_loss:.4f} ftrain_acc{tr_acc:.4f} val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc wait 0 torch.save({ model: model.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch, }, best_wheat_seed.pt) else: wait 1 if wait patience: print(early stopping) breaktorch.no_grad()是评估函数的关键它关闭了自动求图否则验证阶段也会累积中间变量显存和内存都会被慢慢吃光。model.eval()会把 Dropout 关掉、让 BatchNorm 使用运行均值这两点漏一个验证集上的数值都会失真。保存的是验证集准确率最高的那一版权重不是最后一轮权重这样才能避开过拟合末期的低质量快照。4.4 解冻尾部参数做第二轮微调冻结阶段跑完后最佳权重通常在 85% 到 92% 之间。如果还想往上提就把骨干的最后几个 stage 解冻用更小的学习率继续训练ckpt torch.load(best_wheat_seed.pt) model.load_state_dict(ckpt[model]) for name, param in model.named_parameters(): if layer4 in name or layer3 in name or fc in name: param.requires_grad True else: param.requires_grad False optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr1e-4, weight_decay1e-4) # 继续跑同样的训练循环max_epochs 这里给 25解冻层数不要一次给太多layer4 layer3 fc是常见的平衡点。学习率从3e-4降到1e-4否则预训练特征会被大幅度改写。微调阶段每轮训练时间大概是冻结阶段的三倍但准确率的提升通常有 1 到 3 个百分点。5. 避坑2000张小样本图像分类的 5 个常见翻车现场5.1 验证准确率比训练高不少现象跑了几个 epochval_acc一直比train_acc高看起来非常体面。原因最常见的是划分阶段把同源照片随机分到了两边验证集等于在“背答案”另一种可能是验证集继承了一部分训练增强导致训练更难学而验证更简单。解决先回到第 2 章确认是否按样本组划分再检查val_tf和train_tf是不是独立对象别在调试时不小心把变量引到同一个 transform 上最后在验证函数里打印model.training确认处于False状态。5.2 loss 一直降但准确率不动现象train_loss从 2.2 降到 0.9但train_acc始终在 0.4 附近晃明明在收敛却分不出类别。原因类别极度不平衡或者有几个品种从视觉上非常接近。模型先学会的是占主导的信号比如背景亮度而不是品种之间的微细差异。解决别只看总准确率去打印分品种准确率几行代码就能看出是哪几个类互相吞。给CrossEntropyLoss加类别权重是最直接的手段class_weights torch.tensor([1.0, 2.0, 1.0, 1.5, 1.0]).to(device) criterion torch.nn.CrossEntropyLoss(weightclass_weights)权重按各类样本数的反比设置比如样本数少的类别给更高权重。如果还不行再考虑 Focal Loss它会让模型把注意力更多放在难分样本上但先试类别权重别一上来就上复杂方案。5.3 训练集上九成五换一批光照新图立刻掉到六成现象自己拍的验证集上 95% 准确率跑到另外一个产地、换个光源或背景板准确率直接崩。原因模型学的不是种子本身的纹理而是采集环境的恒定特征比如白色背景、固定色温、固定角度。这是小样本图像分类里最普遍过拟合形式。解决数据增强里把ColorJitter的亮度范围开到0.2~0.3对比度也给一定扰动让网络不能依赖固定的颜色统计如果背景是可控制的原始采集时收一批不同背景的照片比增强更治本。5.4 迁移学习一开始就出 NaN现象刚跑第一个 batchloss 飙到 200 多然后变成nan。原因分类头是随机初始化的输出 logits 没有任何限制碰上学习率过大就炸了也可能是数据里有损坏图被convert(RGB)强行转换后变成了异常通道值。解决先把学习率调到3e-5验证模型能正常迭代再逐步放大。同时检查x.dtype是不是torch.float32以及数据里是否存在单通道灰度图被强行复制成三通道的情况——这类样本在ToTensor之后的像素分布会偏离正常范围。5.5 显存不够多半不是数据量大而是验证没关梯度现象2000张图、ResNet50 冻结骨干按理说显存占用很小结果训练到一半 OOM。原因最常见的是自己写的验证函数里漏了torch.no_grad()或者验证循环里没有model.eval()导致 BatchNorm 统计还在更新还保留了大量中间激活。另一类原因调试代码里用了plt.show()但没关窗口GUI 进程积累吃掉内存。解决用上面给出的evaluate函数结构所有推理路径都强制torch.no_grad()DataLoader 的num_workers设到 2 到 4别拉满如果连续跑多轮实验在每轮之间调一次torch.cuda.empty_cache()但不要在每个 step 里调那只会让训练更慢。6. 验证与落地混淆矩阵、可视化与一个保存模型的小习惯6.1 用混淆矩阵找出互相吞并的品种整体准确率只能告诉你模型强不强不能告诉你哪个品种在拖后腿。拿最优权重在验证集上完整跑一遍输出混淆矩阵import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix import numpy as np preds, y_true [], [] model.eval() with torch.no_grad(): for x, y in val_loader: x x.to(device) preds.extend(model(x).argmax(1).cpu().numpy()) y_true.extend(y.numpy()) cm confusion_matrix(y_true, preds, labelslist(range(num_classes))) fig, ax plt.subplots(figsize(6, 5)) im ax.imshow(cm, cmapBlues) for i in range(cm.shape[0]): for j in range(cm.shape[1]): ax.text(j, i, cm[i, j], hacenter, vacenter, colorwhite if cm[i, j] cm.max() / 2 else black) ax.set_xticks(range(num_classes)) ax.set_yticks(range(num_classes)) ax.set_xticklabels(class_names, rotation45) ax.set_yticklabels(class_names) plt.colorbar(im) ax.set_xlabel(预测) ax.set_ylabel(真实) plt.savefig(wheat_confusion.png, dpi150, bbox_inchestight)如果混淆矩阵里某一列的条纹特别多说明模型把多个品种都判成了那一类返回第 5 章用类别权重或增强再调。6.2 用 Grad-CAM 说服别人模型没有抄背景种子分类模型在落地时别人问的第一句话往往是“它到底在看什么”。Grad-CAM 是最直观的验证手段对最后一层卷积的特征图求类别得分的梯度用全局平均得到每个通道的权重再叠加回原图。热力点如果集中在种子腹部沟槽、表皮纹理上那说明模型学到了正确特征如果热力点集中在背景边缘就必须回去修背景和增强。pytorch-grad-cam这个库可以直接用核心逻辑不需要自己从头写。每次验证完我至少会随机挑 10 张图看热力图这比任何准确率指标更能提前暴露采集端的环境泄漏。6.3 把预处理参数和模型文件放在同一目录保存最后一个容易被忽略的细节是部署时的预处理不一致。训练阶段用的是Resize(256) RandomCrop(224)推理阶段往往直接Resize((224, 224))这还可以接受但mean、std和通道顺序一旦没对齐再好的模型权重也会被打回原形。我现在每次训练结束都会把预处理参数存成 JSON 放在模型旁边{ input_size: [224, 224], resize: 256, mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225], normalize: true, channels: 3 }推理脚本统一从这个 JSON 读取参数而不是每次手写一遍 transform。这个习惯是我被一个灰度图样本坑过之后养成的——当时模型训练和验证都很好部署时突然遇到一批 RGB 通道顺序写反的输入图准确率直接掉到随机水平。从那以后模型权重、预处理参数、类别名三件套必须放在同一个目录。希望帮到你。本文还有配套的精品资源点击获取
