尧图网页设计 
简介面向图像分类任务的小麦种子图像分类数据集包含约2000张已标注并完成预处理的图像可直接作为分类网络的输入省去自行清洗与格式转换的步骤也方便与PyTorch、TensorFlow等主流框架对接。资源包共2000个文件以1998张jpg图片为主体另有1个可视化脚本和1个json配置文件压缩后约225.64MB数据已明确划分为训练集与测试集同类图片分目录存放便于按批次加载与模型评估。json文件记录了4个分类类别的详细配置运行show脚本可快速可视化数据集直观检查图片与标签的对应关系提高数据使用效率。目前已有223人学习下载适合计算机视觉初学者用作物种分类练习、算法效果对比也可用于图像分类相关的课程设计或论文实验。配套的图像分类与分割网络改进、计算机视觉完整项目等专栏为后续模型结构调整、迁移学习与工程化部署提供了延伸学习资料。1. 小麦种子图像分类数据集2000张已标注图能跑通怎样的图像分类模型小麦种子图像分类数据集听起来规模很小约2,000张已标注图片连ImageNet的零头都不到。但做过育种选型、粮库复检、种子纯度抽检的工程师会告诉你真实项目里最常遇到的就是这种“够用但不海量”的数据类别能数清楚单类样本从几十到几百张。比数据量更关键的是这批数据怎么组织、标注质量如何、选择什么图像分类算法接得住。这篇笔记按我做过的小样本农作物分类管线来讲先盘数据再写迁移学习训练脚本调参数最后用混淆矩阵和热力图验货。适合手里有类似小数据集、想走通分类全流程的工程师和学生也方便你判断“这类已标注数据集到底够不够投入”。2. 先盘数据目录结构、标签分布与读取脚本2.1 图像分类数据集的“已标注”通常体现在文件夹名上图像分类任务里的“已标注”和检测任务不同一般不给你一个单独的JSON或XML标注文件而是靠文件夹名承载标签也就是每个品种一个小麦种子的类目文件夹里面放着对应图片。要用PyTorch读这批数据最省事的做法是把它整理成ImageFolder能直接吃进去的结构train/和val/下各建若干个品种文件夹文件夹名就是类别名。wheat_dataset/ ├── train/ │ ├── jimai22/ │ │ ├── 0001.jpg │ │ ├── 0002.jpg │ │ └── ... │ ├── bainong207/ │ │ ├── 0001.jpg │ │ └── ... │ └── zhengmai9023/ │ ├── 0001.jpg │ └── ... └── val/ ├── jimai22/ ├── bainong207/ └── zhengmai9023/这里有个容易忽略的点ImageFolder是按文件夹名的字母序生成标签编号的和目录里文件的物理排列顺序无关。也就是说bainong207会拿到 label 0jimai22拿到 label 1zhengmai9023拿到 label 2。如果你后面要输出中文品种名最好在建 dataset 之后立刻读取dataset.class_to_idx存一份映射避免训练完发现推理阶段标签对不上。另外要留意文件夹名千万不能带空格和中文。并不是说 PyTorch 读不了而是后面跨平台传数据、写 CSV 报告、做 Grad-CAM 可视化时文件名里的非 ASCII 字符很容易在各种环境里翻车。数据到位后第一件事不是写网络而是把一个可靠的class_to_idx映射固化下来。2.2 用一段短脚本摸清2000张图的真实家底很多初学者拿到数据集后第一件事是去看别人怎么训练我却建议你先花二十分钟写脚本盘数据。2000张图说多不多如果六个品种样本数差异悬殊比如有的类有400张有的类只有120张后面训练模型一定会发生“偏科”。所以我在做任何分类网络之前都会先统计每个文件夹里的图片数量、图片尺寸和文件格式。import os from PIL import Image from collections import defaultdict root wheat_dataset/train stats defaultdict(lambda: {count: 0, sizes: set(), formats: set()}) for class_name in sorted(os.listdir(root)): class_path os.path.join(root, class_name) if not os.path.isdir(class_path): continue for fn in os.listdir(class_path): if not fn.lower().endswith((.jpg, .jpeg, .png, .bmp)): print(f[警告] 非图片文件: {class_name}/{fn}) continue fpath os.path.join(class_path, fn) with Image.open(fpath) as im: fmt im.format w, h im.size stats[class_name][count] 1 stats[class_name][sizes].add((w, h)) stats[class_name][formats].add(fmt) for cls, s in sorted(stats.items()): print(f{cls}: {s[count]} 张, 尺寸集合{list(s[sizes])[:5]}, 格式{s[formats]}) total sum(s[count] for s in stats.values()) print(f\n总图片数: {total})这段脚本背后有三个动机第一确认每个类的最少样本数如果有个类只有30张按七三开划分后训练集只剩21张这种类基本只能靠数据增强硬撑或者干脆合并到相近的品种里第二确认图片尺寸是否统一小麦种子特写图里常出现个别接近正方形、个别是细长条的情况Resize的参数得按最小可接受尺寸来设计而不是拍脑袋定一个 224×224第三扫描非图片文件。数据集压缩包里经常混入Thumbs.db、.DS_Store有的标注者图省事把 RAW 格式也塞了进来这类文件在 DataLoader 里是隐患。输出结果里如果发现大量 PNG 和 JPEG 混用问题不大Image.open在读取后会自动转成 RGB 三通道张量。但如果出现灰度图就需要注意灰度图Image.open得到的是单通道直接ToTensor之后模型输入通道是 1和预训练模型要求的 3 通道对不上。我通常会在读取分支里加一句im im.convert(RGB)把所有图片统一成三通道。2.3 拿到这批数据后我做的三件事盘完脚本还要做三件事才敢开始训练。第一把每个类别按样本数从少到多排个序最小的类会决定你数据增强的下限。第二随机挑几十张图肉眼过一遍重点看有没有模糊图、严重过曝图、和背景混成一团的图。小麦种子密集排列时如果标注框里同时出现了半粒麦子和杂质这类样本会直接干扰训练。第三确认验证集的划分方式。2000张图不必追求太复杂的交叉验证但必须保证按类别比例切分否则某个类全部落在验证集里训练时这个类就成了黑匣子。我的习惯是先做一轮硬清洗把肉眼可见的模糊图移到reject/文件夹而不是删除这样后面做数据增强和模型分析时还能随时对比“洗前洗后”的效果。这类清洗脚本不复杂但要记录清楚清洗规则比如按像素方差、按长宽比过滤将来写论文报告时能交代得清。3. PyTorch 迁移学习加载这批数据裁剪、增强与归一化参数一次设对3.1 用 ImageFolder 读文件夹标签按字母序自动生成数据盘完就可以加载了。torchvision.datasets.ImageFolder是这类文件夹式分类数据集的默认选择它会把每个子文件夹当成一个类自动扫描图片并且生成dataset.samples和dataset.targets。我通常会配合DataLoader一起用但第一次跑通时建议先单测一下 dataset 的取值。from torchvision import transforms, datasets train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_dataset datasets.ImageFolder(wheat_dataset/train, transformtrain_transform) print(train_dataset.class_to_idx) print(len(train_dataset)) img, label train_dataset[0] print(img.shape, label)这段代码里有两个值得解释的参数。Resize((256, 256))先把短边拉到256CenterCrop(224)再从中间裁出224这套组合是 ImageNet 预训练模型的标准预处理能保留中心区域的种子纹理。Normalize用的 mean 和 std 是 ImageNet 的统计值。如果训练数据集的整体色调明显偏暗黄你可以用自己数据的均值和标准差去替换但常见做法是先用 ImageNet 的统计值跑出一版基线效果不满意时再回算数据集的真实均值和标准差。ImageFolder有个特性图片后缀不写全可能导致漏读默认只认.jpg,.jpeg,.png,.ppm,.bmp,.pgm,.tif,.tiff,.webp如果你有.jfif这类小众后缀lsit 里没被扫到是正常的。第一次加载后先看len(train_dataset)和磁盘里的实际图片数量是否一致不一致就回头盘数据。3.2 按类别比例切训练/验证集别让一个小品类被切空拿到两个 ImageFolder 之前还有一个问题要先定是直接用别人分好的train/val还是自己重新划分。如果发布者已经按类别比例切好了直接用就好如果只给了一个总目录就按类别比例切绝不能直接random_split整个文件列表。直接随机切在小数据集上很容易出现某个类在验证集里占了一半、在训练集里只剩三分之一的情况最后评估出的准确率完全没有参考意义。import os import shutil import numpy as np src_root wheat_dataset/all dst_root wheat_dataset/split val_ratio 0.2 rng np.random.default_rng(42) for class_name in sorted(os.listdir(src_root)): class_dir os.path.join(src_root, class_name) if not os.path.isdir(class_dir): continue files [f for f in sorted(os.listdir(class_dir)) if f.lower().endswith((.jpg, .jpeg, .png))] if len(files) 5: print(f[提示] {class_name} 样本太少无法划分) continue perm rng.permutation(len(files)) n_val max(1, int(len(files) * val_ratio)) val_idx, train_idx perm[:n_val], perm[n_val:] for split, idxs in [(val, val_idx), (train, train_idx)]: out_dir os.path.join(dst_root, split, class_name) os.makedirs(out_dir, exist_okTrue) for i in idxs: src os.path.join(class_dir, files[i]) dst os.path.join(out_dir, files[i]) shutil.copy2(src, dst)这里用shutil.copy2而不是os.rename原因是保留原始数据拆分副本后面如果发现划分有偏见还能迅速重来。rng.random固定种子为42保证每次划分结果一致这一点在复现实验和写论文时非常关键。val_ratio 0.2是默认参数但要注意最小类只有30张时20%的验证集只有6张评估结果的置信度很低。遇到极小的类我会单独把这类的验证比例下调到10%宁可训练时多给几张也不能让验证集里只有一个样本。3.3 训练和验证用两套 transform参数这样设很多入门代码在train_dataset和val_dataset里用同一个 transform这在小数据集的迁移学习里等于放弃了廉价的数据增强。种子图像的拍摄角度、光照、种皮纹理变化很大而同一品种的麦粒又高度相似模型很容易只记住某个特定光源下的颜色分布。训练集必须做随机增强验证集和测试集则保持固定预处理这样才能评估模型在真实分布上的表现。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0), ratio(0.8, 1.2)), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.4, contrast0.4, saturation0.3, hue0.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, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])RandomResizedCrop的scale(0.6, 1.0)是调参重点。默认参数通常是从ImageNet继承来的(0.08, 1.0)但那套参数给目标占据画面大部分的大图设计。小麦种子特写中种子往往已经占满画面中央裁到0.08会让模型经常只看到一粒种子的一部分物理含义不对。所以我限定最小保留60%的区域让随机裁剪更多起到平移、尺度抖动的作用而不是把种子从画面里裁没。ColorJitter的亮度抖动用0.4对比度也用0.4这是拍种子图时最常变化的两个因素。饱和度抖动幅度控制在0.3色相只给0.05因为小麦种子的颜色在红褐色区间色相变化太大会把红皮小麦变成青粒产生物理上不存在的样本。旋转15度对种子足够大角度旋转会让种子的长轴方向发生不符合拍摄习惯的变化反而增加训练难度。4. 训练收敛的开关预训练模型选择、冻结策略与类别不均衡处理4.1 为什么这种规模首选 ResNet-18 而不是直接上 ViT只有2000张数据时最大的风险不是模型不够强而是模型太强导致严重过拟合。视觉Transformer这类需要海量数据才能发挥特性的架构在小数据集上往往不如一个先用ImageNet预训练过的ResNet-18可靠。ResNet-18的backbone只有8个残差块参数量约1100万分类头换成自己的全连接层但在ImageNet上已经学会的纹理、边缘和颜色特征可以直接迁移到麦粒表面纹理识别上。EfficientNet-B0是另一个可选项但要留意它的输入分辨率、缩放系数和训练策略。相比EfficientNetResNet-18在torchvision里自带与版本匹配的预训练权重加载代码最不容易踩坑而且后续做 Grad-CAM 时层名简单不需要去翻复杂的 stage 结构。在这个规模下模型选型的第一原则是“容易复现、容易排错、容易验证”ResNet-18恰好都满足。4.2 冻结backbone只训分类头得到一个能用的基线迁移学习的常规做法是冻结backbone只替换最后的全连接分类头。因为小麦种子图和ImageNet的自然图像共享底层视觉特征比如边缘、颗粒感、光照变化这些不需要重新学。只训练分类头能让模型快速收敛到合理水平同时也防止小数据集在训练初期就把预训练特征带偏。import torch import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) for param in model.parameters(): param.requires_grad False num_classes len(train_dataset.classes) in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes), )这里nn.Dropout(0.3)是给分类头加正则。虽然backbone被冻结但新加的线性层依然有几十万个参数要学2000张数据照样可能在这些层里发生过拟合。Dropout比例0.3比较适合小型数据集太高会让收敛变慢太低起不到抑制作用。in_features在ResNet-18上是512不需要硬编码从model.fc.in_features取最稳妥。训练循环里我习惯用 Adam 优化器只把requires_gradTrue的参数传进去fastai、pytorch-image-models 里的迁移学习教程也大多是这套结构。from torch.optim import Adam from torch.optim.lr_scheduler import ReduceLROnPlateau optimizer Adam(filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4) scheduler ReduceLROnPlateau(optimizer, modemax, factor0.5, patience3, verboseTrue) for epoch in range(30): model.train() train_loss, train_correct, train_total 0.0, 0, 0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() outputs model(images) loss nn.CrossEntropyLoss()(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) train_correct (preds labels).sum().item() train_total labels.size(0) val_acc evaluate(model, val_loader) scheduler.step(val_acc) print(fepoch{epoch1}, loss{train_loss/train_total:.4f}, acc{train_correct/train_total:.4f}, val_acc{val_acc:.4f})lr3e-4对只训练分类头的任务通常是一个安全的起点。如果训练集很小可以再降到1e-4。ReduceLROnPlateau的patience3表示验证集准确率连续三轮不创新高时学习率减半这是避免在后期震荡过大的通用策略。注意weight_decay1e-4虽然加了L2正则但只作用在冻结参数之外的部分所以数值不必设得过大。4.3 类别不均衡时用 WeightedRandomSampler而不是复制样本种子里有的品种常见有的品种稀有单类样本数差三四倍很常见。如果不做处理模型会把概率几乎都压在大类上。我的做法是用WeightedRandomSampler给样本少的类更高被抽中的概率相当于在采样层面做了均衡并不会增加新的图片数据。from torch.utils.data import WeightedRandomSampler targets train_dataset.targets class_counts torch.bincount(torch.tensor(targets)) weights 1.0 / class_counts[targets].float() sampler WeightedRandomSampler(weights, num_sampleslen(targets), replacementTrue) train_loader DataLoader(train_dataset, batch_size16, samplersampler, num_workers2)targets是train_dataset.targets长度和图片数一致。weights的计算方式是先统计每个类别样本数再让每个样本的权重等于所属类别样本数的倒数样本少的类单样本权重更大。num_sampleslen(targets)表示每个epoch 采样的总数和原始样本数一致说明每个epoch 里样本少的小类会被重复抽到而大类会少抽一些。replacementTrue表示每次抽样相互独立同一个图片可以在这个epoch 里出现多次这正是我们要的效果。这里有个容易忽略的坑class_counts[targets]返回的是一个同形状的一维张量里面每个位置是对应图片所属类别的样本数PyTorch 会用广播自动做索引。如果你把targets换成了dataset.samples里提取的类别索引别忘了先torch.tensor再取下标。4.4 什么时候解冻微调看验证曲线别一上来就全量训练冻结backbone训练到一定轮次后验证集准确率会进入平台期这个平台通常在80%到90%之间取决于品种差异和图像质量。此时再继续只训分类头提升空间已经很小。我的经验是先让冻结阶段跑到“连续五轮准确率不再上升”再解冻backbone的最后两层做微调。解冻过早会导致模型重新适应小麦种子数据时把过早的学习率作用在底层特征上破坏预训练特征。解冻时优化器和学习率都必须换。此时不用 Adam 的3e-4而要用更低的学习率常见做法是给param_groups分组backbone部分用1e-5分类头部分用1e-4。如果只有一个优化器就直接把整个模型的学习率降到3e-5。这样再做10~20轮微调通常能再提升2到5个百分点。for name, param in model.named_parameters(): if name.startswith(layer4): param.requires_grad True elif name.startswith(fc): param.requires_grad True这段代码只解冻了layer4和分类头。layer4是ResNet-18最后一个残差块负责最抽象的语义特征离种子品种这种细粒度特征最近先解冻它对预训练知识的破坏最小。如果解冻layer3验证集并不一定会更好在小数据集上反而容易让损失函数在几个epoch后开始震荡。5. 避坑排查2000张种子图上最容易翻车的五个问题5.1 训练loss一直降、验证集acc却原地转圈过拟合的典型信号现象训练集准确率一路涨到99%验证集准确率停留在75%左右而且波动剧烈。我用上面那套冻结backbone的训练方式也出现过一次当时只训了5个epoch就开始在验证集上反复横跳。原因分类头是全连接的线性层加ReLU参数空间足够大2000张图很快就能把训练集背下来。冻结backbone只降低了一部分过拟合风险但没有完全消除。解决先把分类头里的神经元数量从256减到128同时把Dropout提高到0.5学习率从3e-4调低到1e-4。如果还压不住就检查是不是验证集划分时某个品种全部进了训练集另一个品种全部进了验证集这种情况换任何模型都不会收敛。5.2 验证集里某个类分数特别低翻看图片才发现一堆模糊样本现象整体准确率在85%但某个品种的召回率只有60%而且损失函数在验证这个类时明显偏高。原因我最初按文件名统计数量没注意图片清晰度。小麦种子在运输和拍摄过程中难免有虚焦这一类恰好混入了大量模糊图。解决用OpenCV的Laplacian算子计算每个灰度图的方差方差低于20的标记为模糊候选人工确认后移到reject目录。清洗后再训练同一类的准确率通常能提升10个百分点以上。清洗时不要直接删除统一归档方便回溯为什么这个类表现差。5.3 文件夹名带中文或空格DataLoader 报错翻车现象ImageFolder能正常读取目录但DataLoader跑到一半抛异常提示文件路径找不到。原因数据集发布者在整理时用了“济麦22”这类中文名或者在文件名里加了空格。Windows下的路径编码和Linux下的UTF-8处理不一致导致在多进程读取时索引失效。解决将目录统一改成ASCII别名如jimai22用dataset.class_to_idx建立中文名映射。如果不想改目录名就固定num_workers0绕过多进程但训练速度会慢不少。我的习惯是提前在数据清洗环节就把所有非ASCII字符过滤掉。5.4 训练损失下降得很快但验证集出现“同类不同色”的样本模型始终分错现象两个品种颜色非常接近比如红皮小麦和陈年麦粒模型在两者之间反复误判混淆矩阵里它们几乎成了一对一互混的关系。原因图像分类算法只看到了表面颜色和纹理而两个品种在形态上的差异可能集中在腹沟深浅、籽粒长宽比这些像素级特征上普通的数据增强没有放大这些差异。解决如果有已知的品种差异部位可以在预处理里加一个专门裁剪腹沟区域的预处理分支如果没有就调整ColorJitter的色相范围再做一次T-SNE可视化看模型输出的特征向量是否把这两个品种分成了两个明显的簇。如果特征混合在一起说明当前的分辨率不够试一下RandomResizedCrop的scale下限调整为0.5让模型看到更大的种子局部纹理。5.5 盲目把ResNet-18换成ResNet-50性能和训练速度一起变差现象有人看到准确率不够高第一反应是上更大的模型结果ResNet-50验证集准确率反而比ResNet-18低了2个百分点训练时间翻倍。原因更大的模型容量在2000张数据上更容易过拟合而且torchvision里ResNet-50的预训练特征并不比ResNet-18更适合细粒度分类。解决这个规模下先把ResNet-18的微调和数据增强参数调到极限再考虑换模型。如果一定要尝试更强的模型优先试efficientnet_b0或convnext_tiny但必须同时加大数据增强强度和解冻后的正则。模型参数量不是越多越好小数据集的瓶颈在数据分布不在网络容量。6. 用混淆矩阵和 Grad-CAM 给模型“验货”最后导出 ONNX6.1 混淆矩阵比整体 acc 更能反映品种之间的混叠整体准确率再高也无法暴露“哪些品种互相像”这个核心问题。我会在验证集上跑一遍预测把所有(真实标签, 预测标签)记录下来用混淆矩阵观察哪里发生系统性误判。对小麦种子数据重点关注的是对角线之外的高频格子比如jimai22被误判成bainong207说明这两个品种在模型看来特征距离太近。import numpy as np from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.cuda() outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) disp ConfusionMatrixDisplay(confusion_matrixcm) disp.plot(cmapBlues) plt.savefig(confusion_matrix.png, dpi200, bbox_inchestight)这里torch.max(outputs, 1)返回最大值和对应索引取索引就是预测类别。confusion_matrix的第一个参数是真实标签第二个是预测标签别填反了填反的话输出的矩阵是转置的阅读起来会得出错误结论。如果验证集里某个类别只有几十张混淆矩阵的格子数值稀疏可以额外打印classification_report把 Precision/Recall/F1 一起看。6.2 Grad-CAM 看模型到底在盯着种子哪个部位准确率达标后我还会用 Grad-CAM 生成热力图确认模型没有“作弊”。一个典型的翻车案例是模型靠背景中的标签纸颜色分类而不是靠种子形态。热力图会显示模型关注区域集中在种子之外的背景上。如果发现这种问题需要在预处理里做背景裁剪或马赛克而不是继续调模型参数。from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget cam GradCAM(modelmodel, target_layers[model.layer4[-1]]) for images, labels in val_loader: images images.cuda() targets [ClassifierOutputTarget(int(labels[0].item()))] grayscale_cam cam(input_tensorimages[:1], targetstargets) breaktarget_layers选model.layer4[-1]也就是最后一层残差块输出的特征图语义最抽象能反映模型最终的决策依据。ClassifierOutputTarget指定要解释的类别如果传入错误类别热力图反映的内容就不对。这里只取了images[:1]因为 Grad-CAM 对单张图解释最有意义批量解释反而看不出差异。6.3 用 ONNX 导出模型方便脱离 PyTorch 环境推理模型调好后我习惯导出 ONNX。这样做一方面是可以脱离显卡和PyTorch环境用 ONNX Runtime 在普通CPU服务器上做批量推理另一方面也让模型格式保持中立将来换推理框架不用重新训练。导出代码非常简单但有几个参数容易踩坑。model.eval() dummy_input torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, wheat_classifier.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version12, )opset_version12覆盖了绝大多数CPU部署场景过高的版本号在老设备上可能不被支持。dynamic_axes设置batch维度为动态这样才能用一个模型文件处理任意数量的输入图片。导出前必须调用model.eval()否则模型里的Dropout层会以训练模式运行导出的模型推理结果不稳定。导出后我通常先对比 PyTorch 模型和 ONNX 模型在同一张图上的输出差异差异超过 1e-4 就检查一下 batch normalization 层是否被错误冻结。这批数据集只有约2000张图片但足够走通完整的图像分类迁移学习流程。我踩过最大的坑是忽略分类目标里两个品种在腹沟纹理上的差异盲目调网络结构最后发现调整裁剪尺度和光源颜色抖动才是真正的突破口。希望这份从盘数据到导出的完整路径能帮到你也别忘了对这类小数据集先怀疑数据再怀疑模型。本文还有配套的精品资源点击获取 小麦种子图像分类数据集2000张，迁移学习实战全流程指南
简介面向图像分类任务的小麦种子图像分类数据集包含约2000张已标注并完成预处理的图像可直接作为分类网络的输入省去自行清洗与格式转换的步骤也方便与PyTorch、TensorFlow等主流框架对接。资源包共2000个文件以1998张jpg图片为主体另有1个可视化脚本和1个json配置文件压缩后约225.64MB数据已明确划分为训练集与测试集同类图片分目录存放便于按批次加载与模型评估。json文件记录了4个分类类别的详细配置运行show脚本可快速可视化数据集直观检查图片与标签的对应关系提高数据使用效率。目前已有223人学习下载适合计算机视觉初学者用作物种分类练习、算法效果对比也可用于图像分类相关的课程设计或论文实验。配套的图像分类与分割网络改进、计算机视觉完整项目等专栏为后续模型结构调整、迁移学习与工程化部署提供了延伸学习资料。1. 小麦种子图像分类数据集2000张已标注图能跑通怎样的图像分类模型小麦种子图像分类数据集听起来规模很小约2,000张已标注图片连ImageNet的零头都不到。但做过育种选型、粮库复检、种子纯度抽检的工程师会告诉你真实项目里最常遇到的就是这种“够用但不海量”的数据类别能数清楚单类样本从几十到几百张。比数据量更关键的是这批数据怎么组织、标注质量如何、选择什么图像分类算法接得住。这篇笔记按我做过的小样本农作物分类管线来讲先盘数据再写迁移学习训练脚本调参数最后用混淆矩阵和热力图验货。适合手里有类似小数据集、想走通分类全流程的工程师和学生也方便你判断“这类已标注数据集到底够不够投入”。2. 先盘数据目录结构、标签分布与读取脚本2.1 图像分类数据集的“已标注”通常体现在文件夹名上图像分类任务里的“已标注”和检测任务不同一般不给你一个单独的JSON或XML标注文件而是靠文件夹名承载标签也就是每个品种一个小麦种子的类目文件夹里面放着对应图片。要用PyTorch读这批数据最省事的做法是把它整理成ImageFolder能直接吃进去的结构train/和val/下各建若干个品种文件夹文件夹名就是类别名。wheat_dataset/ ├── train/ │ ├── jimai22/ │ │ ├── 0001.jpg │ │ ├── 0002.jpg │ │ └── ... │ ├── bainong207/ │ │ ├── 0001.jpg │ │ └── ... │ └── zhengmai9023/ │ ├── 0001.jpg │ └── ... └── val/ ├── jimai22/ ├── bainong207/ └── zhengmai9023/这里有个容易忽略的点ImageFolder是按文件夹名的字母序生成标签编号的和目录里文件的物理排列顺序无关。也就是说bainong207会拿到 label 0jimai22拿到 label 1zhengmai9023拿到 label 2。如果你后面要输出中文品种名最好在建 dataset 之后立刻读取dataset.class_to_idx存一份映射避免训练完发现推理阶段标签对不上。另外要留意文件夹名千万不能带空格和中文。并不是说 PyTorch 读不了而是后面跨平台传数据、写 CSV 报告、做 Grad-CAM 可视化时文件名里的非 ASCII 字符很容易在各种环境里翻车。数据到位后第一件事不是写网络而是把一个可靠的class_to_idx映射固化下来。2.2 用一段短脚本摸清2000张图的真实家底很多初学者拿到数据集后第一件事是去看别人怎么训练我却建议你先花二十分钟写脚本盘数据。2000张图说多不多如果六个品种样本数差异悬殊比如有的类有400张有的类只有120张后面训练模型一定会发生“偏科”。所以我在做任何分类网络之前都会先统计每个文件夹里的图片数量、图片尺寸和文件格式。import os from PIL import Image from collections import defaultdict root wheat_dataset/train stats defaultdict(lambda: {count: 0, sizes: set(), formats: set()}) for class_name in sorted(os.listdir(root)): class_path os.path.join(root, class_name) if not os.path.isdir(class_path): continue for fn in os.listdir(class_path): if not fn.lower().endswith((.jpg, .jpeg, .png, .bmp)): print(f[警告] 非图片文件: {class_name}/{fn}) continue fpath os.path.join(class_path, fn) with Image.open(fpath) as im: fmt im.format w, h im.size stats[class_name][count] 1 stats[class_name][sizes].add((w, h)) stats[class_name][formats].add(fmt) for cls, s in sorted(stats.items()): print(f{cls}: {s[count]} 张, 尺寸集合{list(s[sizes])[:5]}, 格式{s[formats]}) total sum(s[count] for s in stats.values()) print(f\n总图片数: {total})这段脚本背后有三个动机第一确认每个类的最少样本数如果有个类只有30张按七三开划分后训练集只剩21张这种类基本只能靠数据增强硬撑或者干脆合并到相近的品种里第二确认图片尺寸是否统一小麦种子特写图里常出现个别接近正方形、个别是细长条的情况Resize的参数得按最小可接受尺寸来设计而不是拍脑袋定一个 224×224第三扫描非图片文件。数据集压缩包里经常混入Thumbs.db、.DS_Store有的标注者图省事把 RAW 格式也塞了进来这类文件在 DataLoader 里是隐患。输出结果里如果发现大量 PNG 和 JPEG 混用问题不大Image.open在读取后会自动转成 RGB 三通道张量。但如果出现灰度图就需要注意灰度图Image.open得到的是单通道直接ToTensor之后模型输入通道是 1和预训练模型要求的 3 通道对不上。我通常会在读取分支里加一句im im.convert(RGB)把所有图片统一成三通道。2.3 拿到这批数据后我做的三件事盘完脚本还要做三件事才敢开始训练。第一把每个类别按样本数从少到多排个序最小的类会决定你数据增强的下限。第二随机挑几十张图肉眼过一遍重点看有没有模糊图、严重过曝图、和背景混成一团的图。小麦种子密集排列时如果标注框里同时出现了半粒麦子和杂质这类样本会直接干扰训练。第三确认验证集的划分方式。2000张图不必追求太复杂的交叉验证但必须保证按类别比例切分否则某个类全部落在验证集里训练时这个类就成了黑匣子。我的习惯是先做一轮硬清洗把肉眼可见的模糊图移到reject/文件夹而不是删除这样后面做数据增强和模型分析时还能随时对比“洗前洗后”的效果。这类清洗脚本不复杂但要记录清楚清洗规则比如按像素方差、按长宽比过滤将来写论文报告时能交代得清。3. PyTorch 迁移学习加载这批数据裁剪、增强与归一化参数一次设对3.1 用 ImageFolder 读文件夹标签按字母序自动生成数据盘完就可以加载了。torchvision.datasets.ImageFolder是这类文件夹式分类数据集的默认选择它会把每个子文件夹当成一个类自动扫描图片并且生成dataset.samples和dataset.targets。我通常会配合DataLoader一起用但第一次跑通时建议先单测一下 dataset 的取值。from torchvision import transforms, datasets train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_dataset datasets.ImageFolder(wheat_dataset/train, transformtrain_transform) print(train_dataset.class_to_idx) print(len(train_dataset)) img, label train_dataset[0] print(img.shape, label)这段代码里有两个值得解释的参数。Resize((256, 256))先把短边拉到256CenterCrop(224)再从中间裁出224这套组合是 ImageNet 预训练模型的标准预处理能保留中心区域的种子纹理。Normalize用的 mean 和 std 是 ImageNet 的统计值。如果训练数据集的整体色调明显偏暗黄你可以用自己数据的均值和标准差去替换但常见做法是先用 ImageNet 的统计值跑出一版基线效果不满意时再回算数据集的真实均值和标准差。ImageFolder有个特性图片后缀不写全可能导致漏读默认只认.jpg,.jpeg,.png,.ppm,.bmp,.pgm,.tif,.tiff,.webp如果你有.jfif这类小众后缀lsit 里没被扫到是正常的。第一次加载后先看len(train_dataset)和磁盘里的实际图片数量是否一致不一致就回头盘数据。3.2 按类别比例切训练/验证集别让一个小品类被切空拿到两个 ImageFolder 之前还有一个问题要先定是直接用别人分好的train/val还是自己重新划分。如果发布者已经按类别比例切好了直接用就好如果只给了一个总目录就按类别比例切绝不能直接random_split整个文件列表。直接随机切在小数据集上很容易出现某个类在验证集里占了一半、在训练集里只剩三分之一的情况最后评估出的准确率完全没有参考意义。import os import shutil import numpy as np src_root wheat_dataset/all dst_root wheat_dataset/split val_ratio 0.2 rng np.random.default_rng(42) for class_name in sorted(os.listdir(src_root)): class_dir os.path.join(src_root, class_name) if not os.path.isdir(class_dir): continue files [f for f in sorted(os.listdir(class_dir)) if f.lower().endswith((.jpg, .jpeg, .png))] if len(files) 5: print(f[提示] {class_name} 样本太少无法划分) continue perm rng.permutation(len(files)) n_val max(1, int(len(files) * val_ratio)) val_idx, train_idx perm[:n_val], perm[n_val:] for split, idxs in [(val, val_idx), (train, train_idx)]: out_dir os.path.join(dst_root, split, class_name) os.makedirs(out_dir, exist_okTrue) for i in idxs: src os.path.join(class_dir, files[i]) dst os.path.join(out_dir, files[i]) shutil.copy2(src, dst)这里用shutil.copy2而不是os.rename原因是保留原始数据拆分副本后面如果发现划分有偏见还能迅速重来。rng.random固定种子为42保证每次划分结果一致这一点在复现实验和写论文时非常关键。val_ratio 0.2是默认参数但要注意最小类只有30张时20%的验证集只有6张评估结果的置信度很低。遇到极小的类我会单独把这类的验证比例下调到10%宁可训练时多给几张也不能让验证集里只有一个样本。3.3 训练和验证用两套 transform参数这样设很多入门代码在train_dataset和val_dataset里用同一个 transform这在小数据集的迁移学习里等于放弃了廉价的数据增强。种子图像的拍摄角度、光照、种皮纹理变化很大而同一品种的麦粒又高度相似模型很容易只记住某个特定光源下的颜色分布。训练集必须做随机增强验证集和测试集则保持固定预处理这样才能评估模型在真实分布上的表现。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0), ratio(0.8, 1.2)), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.4, contrast0.4, saturation0.3, hue0.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, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])RandomResizedCrop的scale(0.6, 1.0)是调参重点。默认参数通常是从ImageNet继承来的(0.08, 1.0)但那套参数给目标占据画面大部分的大图设计。小麦种子特写中种子往往已经占满画面中央裁到0.08会让模型经常只看到一粒种子的一部分物理含义不对。所以我限定最小保留60%的区域让随机裁剪更多起到平移、尺度抖动的作用而不是把种子从画面里裁没。ColorJitter的亮度抖动用0.4对比度也用0.4这是拍种子图时最常变化的两个因素。饱和度抖动幅度控制在0.3色相只给0.05因为小麦种子的颜色在红褐色区间色相变化太大会把红皮小麦变成青粒产生物理上不存在的样本。旋转15度对种子足够大角度旋转会让种子的长轴方向发生不符合拍摄习惯的变化反而增加训练难度。4. 训练收敛的开关预训练模型选择、冻结策略与类别不均衡处理4.1 为什么这种规模首选 ResNet-18 而不是直接上 ViT只有2000张数据时最大的风险不是模型不够强而是模型太强导致严重过拟合。视觉Transformer这类需要海量数据才能发挥特性的架构在小数据集上往往不如一个先用ImageNet预训练过的ResNet-18可靠。ResNet-18的backbone只有8个残差块参数量约1100万分类头换成自己的全连接层但在ImageNet上已经学会的纹理、边缘和颜色特征可以直接迁移到麦粒表面纹理识别上。EfficientNet-B0是另一个可选项但要留意它的输入分辨率、缩放系数和训练策略。相比EfficientNetResNet-18在torchvision里自带与版本匹配的预训练权重加载代码最不容易踩坑而且后续做 Grad-CAM 时层名简单不需要去翻复杂的 stage 结构。在这个规模下模型选型的第一原则是“容易复现、容易排错、容易验证”ResNet-18恰好都满足。4.2 冻结backbone只训分类头得到一个能用的基线迁移学习的常规做法是冻结backbone只替换最后的全连接分类头。因为小麦种子图和ImageNet的自然图像共享底层视觉特征比如边缘、颗粒感、光照变化这些不需要重新学。只训练分类头能让模型快速收敛到合理水平同时也防止小数据集在训练初期就把预训练特征带偏。import torch import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) for param in model.parameters(): param.requires_grad False num_classes len(train_dataset.classes) in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, 256), nn.ReLU(inplaceTrue), nn.Linear(256, num_classes), )这里nn.Dropout(0.3)是给分类头加正则。虽然backbone被冻结但新加的线性层依然有几十万个参数要学2000张数据照样可能在这些层里发生过拟合。Dropout比例0.3比较适合小型数据集太高会让收敛变慢太低起不到抑制作用。in_features在ResNet-18上是512不需要硬编码从model.fc.in_features取最稳妥。训练循环里我习惯用 Adam 优化器只把requires_gradTrue的参数传进去fastai、pytorch-image-models 里的迁移学习教程也大多是这套结构。from torch.optim import Adam from torch.optim.lr_scheduler import ReduceLROnPlateau optimizer Adam(filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4) scheduler ReduceLROnPlateau(optimizer, modemax, factor0.5, patience3, verboseTrue) for epoch in range(30): model.train() train_loss, train_correct, train_total 0.0, 0, 0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() outputs model(images) loss nn.CrossEntropyLoss()(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) train_correct (preds labels).sum().item() train_total labels.size(0) val_acc evaluate(model, val_loader) scheduler.step(val_acc) print(fepoch{epoch1}, loss{train_loss/train_total:.4f}, acc{train_correct/train_total:.4f}, val_acc{val_acc:.4f})lr3e-4对只训练分类头的任务通常是一个安全的起点。如果训练集很小可以再降到1e-4。ReduceLROnPlateau的patience3表示验证集准确率连续三轮不创新高时学习率减半这是避免在后期震荡过大的通用策略。注意weight_decay1e-4虽然加了L2正则但只作用在冻结参数之外的部分所以数值不必设得过大。4.3 类别不均衡时用 WeightedRandomSampler而不是复制样本种子里有的品种常见有的品种稀有单类样本数差三四倍很常见。如果不做处理模型会把概率几乎都压在大类上。我的做法是用WeightedRandomSampler给样本少的类更高被抽中的概率相当于在采样层面做了均衡并不会增加新的图片数据。from torch.utils.data import WeightedRandomSampler targets train_dataset.targets class_counts torch.bincount(torch.tensor(targets)) weights 1.0 / class_counts[targets].float() sampler WeightedRandomSampler(weights, num_sampleslen(targets), replacementTrue) train_loader DataLoader(train_dataset, batch_size16, samplersampler, num_workers2)targets是train_dataset.targets长度和图片数一致。weights的计算方式是先统计每个类别样本数再让每个样本的权重等于所属类别样本数的倒数样本少的类单样本权重更大。num_sampleslen(targets)表示每个epoch 采样的总数和原始样本数一致说明每个epoch 里样本少的小类会被重复抽到而大类会少抽一些。replacementTrue表示每次抽样相互独立同一个图片可以在这个epoch 里出现多次这正是我们要的效果。这里有个容易忽略的坑class_counts[targets]返回的是一个同形状的一维张量里面每个位置是对应图片所属类别的样本数PyTorch 会用广播自动做索引。如果你把targets换成了dataset.samples里提取的类别索引别忘了先torch.tensor再取下标。4.4 什么时候解冻微调看验证曲线别一上来就全量训练冻结backbone训练到一定轮次后验证集准确率会进入平台期这个平台通常在80%到90%之间取决于品种差异和图像质量。此时再继续只训分类头提升空间已经很小。我的经验是先让冻结阶段跑到“连续五轮准确率不再上升”再解冻backbone的最后两层做微调。解冻过早会导致模型重新适应小麦种子数据时把过早的学习率作用在底层特征上破坏预训练特征。解冻时优化器和学习率都必须换。此时不用 Adam 的3e-4而要用更低的学习率常见做法是给param_groups分组backbone部分用1e-5分类头部分用1e-4。如果只有一个优化器就直接把整个模型的学习率降到3e-5。这样再做10~20轮微调通常能再提升2到5个百分点。for name, param in model.named_parameters(): if name.startswith(layer4): param.requires_grad True elif name.startswith(fc): param.requires_grad True这段代码只解冻了layer4和分类头。layer4是ResNet-18最后一个残差块负责最抽象的语义特征离种子品种这种细粒度特征最近先解冻它对预训练知识的破坏最小。如果解冻layer3验证集并不一定会更好在小数据集上反而容易让损失函数在几个epoch后开始震荡。5. 避坑排查2000张种子图上最容易翻车的五个问题5.1 训练loss一直降、验证集acc却原地转圈过拟合的典型信号现象训练集准确率一路涨到99%验证集准确率停留在75%左右而且波动剧烈。我用上面那套冻结backbone的训练方式也出现过一次当时只训了5个epoch就开始在验证集上反复横跳。原因分类头是全连接的线性层加ReLU参数空间足够大2000张图很快就能把训练集背下来。冻结backbone只降低了一部分过拟合风险但没有完全消除。解决先把分类头里的神经元数量从256减到128同时把Dropout提高到0.5学习率从3e-4调低到1e-4。如果还压不住就检查是不是验证集划分时某个品种全部进了训练集另一个品种全部进了验证集这种情况换任何模型都不会收敛。5.2 验证集里某个类分数特别低翻看图片才发现一堆模糊样本现象整体准确率在85%但某个品种的召回率只有60%而且损失函数在验证这个类时明显偏高。原因我最初按文件名统计数量没注意图片清晰度。小麦种子在运输和拍摄过程中难免有虚焦这一类恰好混入了大量模糊图。解决用OpenCV的Laplacian算子计算每个灰度图的方差方差低于20的标记为模糊候选人工确认后移到reject目录。清洗后再训练同一类的准确率通常能提升10个百分点以上。清洗时不要直接删除统一归档方便回溯为什么这个类表现差。5.3 文件夹名带中文或空格DataLoader 报错翻车现象ImageFolder能正常读取目录但DataLoader跑到一半抛异常提示文件路径找不到。原因数据集发布者在整理时用了“济麦22”这类中文名或者在文件名里加了空格。Windows下的路径编码和Linux下的UTF-8处理不一致导致在多进程读取时索引失效。解决将目录统一改成ASCII别名如jimai22用dataset.class_to_idx建立中文名映射。如果不想改目录名就固定num_workers0绕过多进程但训练速度会慢不少。我的习惯是提前在数据清洗环节就把所有非ASCII字符过滤掉。5.4 训练损失下降得很快但验证集出现“同类不同色”的样本模型始终分错现象两个品种颜色非常接近比如红皮小麦和陈年麦粒模型在两者之间反复误判混淆矩阵里它们几乎成了一对一互混的关系。原因图像分类算法只看到了表面颜色和纹理而两个品种在形态上的差异可能集中在腹沟深浅、籽粒长宽比这些像素级特征上普通的数据增强没有放大这些差异。解决如果有已知的品种差异部位可以在预处理里加一个专门裁剪腹沟区域的预处理分支如果没有就调整ColorJitter的色相范围再做一次T-SNE可视化看模型输出的特征向量是否把这两个品种分成了两个明显的簇。如果特征混合在一起说明当前的分辨率不够试一下RandomResizedCrop的scale下限调整为0.5让模型看到更大的种子局部纹理。5.5 盲目把ResNet-18换成ResNet-50性能和训练速度一起变差现象有人看到准确率不够高第一反应是上更大的模型结果ResNet-50验证集准确率反而比ResNet-18低了2个百分点训练时间翻倍。原因更大的模型容量在2000张数据上更容易过拟合而且torchvision里ResNet-50的预训练特征并不比ResNet-18更适合细粒度分类。解决这个规模下先把ResNet-18的微调和数据增强参数调到极限再考虑换模型。如果一定要尝试更强的模型优先试efficientnet_b0或convnext_tiny但必须同时加大数据增强强度和解冻后的正则。模型参数量不是越多越好小数据集的瓶颈在数据分布不在网络容量。6. 用混淆矩阵和 Grad-CAM 给模型“验货”最后导出 ONNX6.1 混淆矩阵比整体 acc 更能反映品种之间的混叠整体准确率再高也无法暴露“哪些品种互相像”这个核心问题。我会在验证集上跑一遍预测把所有(真实标签, 预测标签)记录下来用混淆矩阵观察哪里发生系统性误判。对小麦种子数据重点关注的是对角线之外的高频格子比如jimai22被误判成bainong207说明这两个品种在模型看来特征距离太近。import numpy as np from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.cuda() outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) disp ConfusionMatrixDisplay(confusion_matrixcm) disp.plot(cmapBlues) plt.savefig(confusion_matrix.png, dpi200, bbox_inchestight)这里torch.max(outputs, 1)返回最大值和对应索引取索引就是预测类别。confusion_matrix的第一个参数是真实标签第二个是预测标签别填反了填反的话输出的矩阵是转置的阅读起来会得出错误结论。如果验证集里某个类别只有几十张混淆矩阵的格子数值稀疏可以额外打印classification_report把 Precision/Recall/F1 一起看。6.2 Grad-CAM 看模型到底在盯着种子哪个部位准确率达标后我还会用 Grad-CAM 生成热力图确认模型没有“作弊”。一个典型的翻车案例是模型靠背景中的标签纸颜色分类而不是靠种子形态。热力图会显示模型关注区域集中在种子之外的背景上。如果发现这种问题需要在预处理里做背景裁剪或马赛克而不是继续调模型参数。from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget cam GradCAM(modelmodel, target_layers[model.layer4[-1]]) for images, labels in val_loader: images images.cuda() targets [ClassifierOutputTarget(int(labels[0].item()))] grayscale_cam cam(input_tensorimages[:1], targetstargets) breaktarget_layers选model.layer4[-1]也就是最后一层残差块输出的特征图语义最抽象能反映模型最终的决策依据。ClassifierOutputTarget指定要解释的类别如果传入错误类别热力图反映的内容就不对。这里只取了images[:1]因为 Grad-CAM 对单张图解释最有意义批量解释反而看不出差异。6.3 用 ONNX 导出模型方便脱离 PyTorch 环境推理模型调好后我习惯导出 ONNX。这样做一方面是可以脱离显卡和PyTorch环境用 ONNX Runtime 在普通CPU服务器上做批量推理另一方面也让模型格式保持中立将来换推理框架不用重新训练。导出代码非常简单但有几个参数容易踩坑。model.eval() dummy_input torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, wheat_classifier.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version12, )opset_version12覆盖了绝大多数CPU部署场景过高的版本号在老设备上可能不被支持。dynamic_axes设置batch维度为动态这样才能用一个模型文件处理任意数量的输入图片。导出前必须调用model.eval()否则模型里的Dropout层会以训练模式运行导出的模型推理结果不稳定。导出后我通常先对比 PyTorch 模型和 ONNX 模型在同一张图上的输出差异差异超过 1e-4 就检查一下 batch normalization 层是否被错误冻结。这批数据集只有约2000张图片但足够走通完整的图像分类迁移学习流程。我踩过最大的坑是忽略分类目标里两个品种在腹沟纹理上的差异盲目调网络结构最后发现调整裁剪尺度和光源颜色抖动才是真正的突破口。希望这份从盘数据到导出的完整路径能帮到你也别忘了对这类小数据集先怀疑数据再怀疑模型。本文还有配套的精品资源点击获取