☰
小样本小麦种子图像分类:迁移学习与ResNet50实战
2026/9/28 1:03:23 网站建设 项目流程

简介:小麦种子图像分类数据集是一份面向图像分类任务学习与研究的已标注数据资源,适合深度学习初学者或农业表型分析人员用于训练和验证分类网络。数据集对原始图片完成了预处理,共划分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(normalize=True))

这里要提醒一句: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, transform=None): 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, label

Dataset本身不做归一化,归一化全部交给transform。这里把通道处理写在数据读取层,比在训练循环里判断通道数要清爽得多。

transform是预处理的关键,也是训练和验证必须分开的地方:

from torchvision import transforms train_tf = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(degrees=15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.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(weights=models.ResNet50_Weights.IMAGENET1K_V1) fc_in = model.fc.in_features model.fc = nn.Sequential( nn.Dropout(p=0.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)

weights=models.ResNet50_Weights.IMAGENET1K_V1是当前 torchvision 的推荐写法,别再用旧版的pretrained=True,那个参数在新版本里被移除了。分类头从单层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_lr=1e-6, max_lr=5e-3, steps=20): 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()), lr=3e-4, weight_decay=1e-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(f"epoch {epoch+1:02d} train_loss={tr_loss:.4f} val_loss={val_loss:.4f} " f"train_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") break

@torch.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()), lr=1e-4, weight_decay=1e-4) # 继续跑同样的训练循环,max_epochs 这里给 25

解冻层数不要一次给太多,layer4 + layer3 + fc是常见的平衡点。学习率从3e-4降到1e-4,否则预训练特征会被大幅度改写。微调阶段每轮训练时间大概是冻结阶段的三倍,但准确率的提升通常有 1 到 3 个百分点。

5. 避坑:2000张小样本图像分类的 5 个常见翻车现场

5.1 验证准确率比训练高不少

现象:跑了几个 epoch,val_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(weight=class_weights)

权重按各类样本数的反比设置,比如样本数少的类别给更高权重。如果还不行,再考虑 Focal Loss,它会让模型把注意力更多放在难分样本上,但先试类别权重,别一上来就上复杂方案。

5.3 训练集上九成五,换一批光照新图立刻掉到六成

现象:自己拍的验证集上 95% 准确率,跑到另外一个产地、换个光源或背景板,准确率直接崩。

原因:模型学的不是种子本身的纹理,而是采集环境的恒定特征,比如白色背景、固定色温、固定角度。这是小样本图像分类里最普遍过拟合形式。解决:数据增强里把ColorJitter的亮度范围开到0.2~0.3,对比度也给一定扰动,让网络不能依赖固定的颜色统计;如果背景是可控制的,原始采集时收一批不同背景的照片比增强更治本。

5.4 迁移学习一开始就出 NaN

现象:刚跑第一个 batch,loss 飙到 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, labels=list(range(num_classes))) fig, ax = plt.subplots(figsize=(6, 5)) im = ax.imshow(cm, cmap="Blues") for i in range(cm.shape[0]): for j in range(cm.shape[1]): ax.text(j, i, cm[i, j], ha="center", va="center", color="white" 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, rotation=45) ax.set_yticklabels(class_names) plt.colorbar(im) ax.set_xlabel("预测") ax.set_ylabel("真实") plt.savefig("wheat_confusion.png", dpi=150, bbox_inches="tight")

如果混淆矩阵里某一列的条纹特别多,说明模型把多个品种都判成了那一类,返回第 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 通道顺序写反的输入图,准确率直接掉到随机水平。从那以后,模型权重、预处理参数、类别名三件套必须放在同一个目录。希望帮到你。

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询