简介面向图像分类学习与算法验证场景这份已标注的马铃薯食物缺陷图像分类数据集提供约450张样本覆盖Dry Rot、Blackleg、Pink Rot等7类缺陷适合初学者快速开展CNN分类训练、迁移学习或数据增强实验也可作为农业图像识别课题的小规模基准数据。包体总文件数为454个以446张jpg主图、5张jpeg补充图像为主另附1个py可视化脚本和1个json分类标签文件整体压缩包约12.58MB结构轻量下载后可直接读取。已有150人学习下载。数据规模适中兼顾类别覆盖与训练效率便于快速建立图像分类基线。数据集已按训练集与测试集划分同一类图片分别存放运行show脚本即可快速预览各缺陷类别配合json文件可核对标签与类别映射关系减少手动整理数据的时间。对于关注CNN分类网络改进的读者还能在此基础上复现实验、对比不同网络结构的分类效果并进一步用于数据增强、模型调参等扩展练习。1. 只有 450 张图的马铃薯缺陷数据集为什么说是小样本分类的练手神器做图像分类最怕什么不是模型结构而是没数据或者数据标得乱七八糟。这份马铃薯食物缺陷图像分类数据集总共约 450 张已标注图片分成 Dry Rot、Blackleg、Pink Rot 等 7 类划分好了训练集和测试集还带一个 show 脚本方便可视化。对想学图像分类算法、想复现 CNN 改进论文、或者准备跑通 ResNet50 小样本图像分类全流程的人来说属于拿来就能用的标准数据。它最大的价值不在于数量而在于“麻雀虽小五脏俱全”能帮你把数据校验、模型训练、分类头调整、模型导出这条链路完整走一遍。2. 先拆数据集再谈训练目录结构、JSON 标签与可视化脚本很多初学者拿到数据集第一件事就是写模型然后开训结果训练半天准确率上不去回头排查才发现是数据读错了。我的习惯是拿到任何数据集先花半小时把目录和标签彻底摸清楚。这一步能帮你避开后面至少一半的“玄学”问题。2.1 目录结构与标注文件不要凭文件名猜类别文件名里的57.jpeg、42.jpeg这种纯粹是图片的编号完全不能代表类别信息。真正的标签都在项目提供的 JSON 文件里也就是说你需要通过 JSON 去拿到每张图片对应的类别字符串再去和训练文件夹里的图片做映射。# 先看整体目录结构确认训练集、测试集和脚本的摆放位置 tree -L 2import json from collections import Counter # 常见做法是json文件里存的是 {图片文件名: 类别名} # 例如 {57.jpeg: Dry Rot, 42.jpeg: Blackleg} with open(label.json, r, encodingutf-8) as f: labels json.load(f) # 统计类别数和每个类别的样本量 class_names sorted(set(labels.values())) print(类别数:, len(class_names)) print(类别列表:, class_names) print(样本分布:, Counter(labels.values()))这段代码帮你确认两件事第一类别数是不是摘要里说的 7 类第二每个类别的样本量是不是均衡。如果发现某个类别只有 20 张图而另一个类别有 120 张这就是明显的类别不平衡需要在训练时考虑加权采样或者 Focal Loss。2.2 运行 show 脚本验证图片与标签是否对得上数据里最经典的坑就是图跟标签错位。图片是 Dry Rot标签写的是 Blackleg。项目自带的 show 脚本就是用来干这件事的它会随机挑几张图把图片和对应的类别标签显示在同一张画布上。# 直接运行项目自带的可视化脚本 python show.py如果脚本默认读取的路径和你解压后的路径不一致通常报FileNotFoundError。这时候别急着改代码先看脚本开头的DATA_DIR xxx这个全局变量把它改成你实际的图片路径即可。# 如果自带脚本跑不起来可以用这个极简版替代可视化 import matplotlib.pyplot as plt import matplotlib.image as mpimg fig, axes plt.subplots(2, 4, figsize(12, 6)) for i, ax in enumerate(axes.flat): img_path list(labels.keys())[i] ax.imshow(mpimg.imread(img_path)) ax.set_title(labels[img_path], fontsize10) ax.axis(off) plt.tight_layout() plt.show()提示可视化不是走形式。如果你发现某个类别的图片里混入了其他类别的背景或者图片是灰度图而代码默认读三通道这类问题会在训练时直接拉低准确率。每张图的标题和图片内容对不上就说明数据本身需要清洗。3. 用 ResNet50 训练基线完整参数设置与小样本迁移学习小样本图像分类的难点在于450 张图根本不够让模型从零开始学习。随机初始化的 ResNet50 在这种数据量下会严重过拟合训练集准确率能到 95%验证集却只有 60%。正确做法是使用在 ImageNet 上预训练过的权重把小样本分类当成一个迁移学习任务来处理。3.1 为什么选 ResNet50小样本下的迁移学习策略选 ResNet50 而不是 ViT 或 ResNet101主要原因是 450 张图的数据量撑不起大模型。ViT 这类模型动辄上亿参数没有几万张图很难发挥出注意力机制的优势ResNet101 比 ResNet50 深得多但在这个数据量下的提升非常有限反而更容易过拟合。ResNet50 在 ImageNet 上学到的边缘、纹理、颜色特征恰好能迁移到马铃薯表面的腐烂区域识别上。对比来看ResNet50 的预训练权重普及度高、显存占用适中224x224 输入下大概 2.5G 显存而且改分类头的成本极低。这就是我把它作为首选基线的原因。3.2 数据增强与超参数配置让 450 张图发挥最大价值数据增强是小样本训练里最实用的“后悔药”。通过在训练时随机裁剪、翻转、调色相当于把 450 张图扩充成了上千个不同视角的变体。但要注意增强强度不能太大否则会把腐烂区域的颜色信息破坏掉反而误导模型。from torchvision import transforms # 小样本图像分类的标准预处理流程 train_transforms transforms.Compose([ transforms.Resize(256), # 先把短边统一到256 transforms.RandomResizedCrop(224, scale(0.7, 1.0)), # 随机裁剪模拟多尺度 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), # 颜色扰动 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 验证集不需要增强只做等比缩放和中心裁剪 val_transforms transforms.Compose([ transforms.Resize(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.7到1.0原因是马铃薯缺陷区域通常占据图片的比例较大如果裁剪比例太小模型会看到一整块正常表皮学不到缺陷特征。ColorJitter的强度也刻意调得比较保守防止真实腐烂颜色被扭曲。训练超参数我建议采用下面的配置这是在小样本分类里比较稳的组合参数推荐值调整建议Batch Size16 或 32数据集小batch 太大容易直接记住训练集OptimizerAdamW权重衰减设 0.01比 SGD 在小数据上收敛更快Learning Rate1e-4预训练模型微调用 1e-4从头训练才用 1e-3Epochs50 或 80配合余弦退火后期学习率极小微调更细致SchedulerCosineAnnealingLR避免学习率陡降导致的 loss 反弹3.3 训练脚本与日志解读把模型包装好之后训练循环本身并不复杂关键是要在每个 epoch 结束后跑验证集并且把 loss 和准确率打出来看趋势。import torch import torch.nn as nn from torchvision import models # 载入ImageNet预训练权重 model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) # 改分类头: 原模型fc层输出1000维这里必须换成7 model.fc nn.Linear(model.fc.in_features, 7) model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.01) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50)从代码里能看到改 ResNet50 只需要动model.fc一行。model.fc.in_features会自动读取原分类头的输入维度即 2048不需要硬编码这样以后换 ResNet18 或 ResNet101 也能通用。训练过程里最需要盯住的信号是验证集准确率如果连续 10 个 epoch 没有上升同时训练集准确率已经逼近 100%那就是典型的过拟合信号。此时优先降低学习率其次增加 Dropout 或增强数据扰动强度而不是盲目堆 epoch。4. 改进基线模型给 ResNet50 加 CBAM 注意力并重构分类头数据集本身附带的信息里提到了 CNN 分类网络改进这也是大多数人下载这个数据集的核心动机——用它来验证自己提出的改进模块有没有效果。我这里给出一个标准的改进思路它能直接在 ResNet50 上提升 2 到 3 个点而且代码量很小。4.1 给 ResNet50 加 CBAM空间与通道注意力CBAM 的原理很好懂就是让网络在学习时更关注“哪里”有缺陷空间注意力和“哪一层特征”对分类最有利通道注意力。马铃薯缺陷分类场景里正常表皮是浅黄色腐烂区域通常是褐色或黑色加注意力机制之后模型会把高权重集中到腐烂区域的纹理边界上而不是被整颗马铃薯的大面积背景带偏。import torch import torch.nn as nn class ChannelAttention(nn.Module): def __init__(self, in_planes, ratio16): super().__init__() self.avg_pool nn.AdaptiveAvgPool2d(1) self.max_pool nn.AdaptiveMaxPool2d(1) self.mlp nn.Sequential( nn.Conv2d(in_planes, in_planes // ratio, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(in_planes // ratio, in_planes, 1, biasFalse) ) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out self.mlp(self.avg_pool(x)) max_out self.mlp(self.max_pool(x)) return self.sigmoid(avg_out max_out) class SpatialAttention(nn.Module): def __init__(self, kernel_size7): super().__init__() self.conv nn.Conv2d(2, 1, kernel_size, paddingkernel_size // 2) self.sigmoid nn.Sigmoid() def forward(self, x): avg_out torch.mean(x, dim1, keepdimTrue) max_out, _ torch.max(x, dim1, keepdimTrue) feature torch.cat([avg_out, max_out], dim1) return self.sigmoid(self.conv(feature)) class CBAM(nn.Module): def __init__(self, in_planes): super().__init__() self.ca ChannelAttention(in_planes) self.sa SpatialAttention() def forward(self, x): x self.ca(x) * x x self.sa(x) * x return x代码里的ratio16控制通道注意力的压缩率压缩越多参数越少但信息损失也越大这个值在 8 到 16 之间效果都比较稳定。空间注意力里的kernel_size7是论文里的默认配置感受野较大能覆盖到马铃薯表面比较大的腐烂区域。插入 CBAM 的常见做法是放在 ResNet50 的layer4之后。因为layer4输出的特征图分辨率最低7x7语义信息最丰富在它后面做注意力加权效率最高。# 在layer4的残差块序列后再挂一个CBAM模块 model.layer4.add_module(cbam, CBAM(2048))4.2 改进分类头从全连接层到全局池化与 Dropout很多人常常只关注主干网络忽略了分类头的重要性。原版 ResNet50 的分类头就是一个简单的全连接层直接输入 2048 个均值池化后的特征。在 450 张图的小样本设定下这个分类头非常容易过拟合。常见做法是在全连接层前加 Dropout并用一个小技巧把原来的池化层显式地暴露出来方便后续做特征导出。from torchvision import models class ImprovedResNet50(nn.Module): def __init__(self, num_classes7): super().__init__() self.backbone models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) self.backbone.fc nn.Identity() # 去掉原来的分类头保留特征提取能力 self.cbam CBAM(2048) self.pool nn.AdaptiveAvgPool2d(1) self.dropout nn.Dropout(0.2) self.fc nn.Linear(2048, num_classes) def forward(self, x): x self.backbone(x) # 输出形状: (batch, 2048, 1, 1) x self.cbam(x) x self.pool(x).flatten(1) x self.dropout(x) return self.fc(x)这里把self.backbone.fc替换成nn.Identity()意味着骨干网络输出的特征图直接输出不再自行展开成向量。然后在外面手动挂pool、dropout和fc。Dropout(0.2)的概率是经验值对小数据集来说0.2 到 0.3 之间比较合适太高会让模型欠拟合。顺便回应一个热搜里常见的疑问用 ViT 评估时分类头需要调整吗答案是必须调整。不管主干是 ResNet 还是 ViT预训练分类头都是为 ImageNet 的 1000 类设计的换到马铃薯 7 分类数据集时分类头的输出维度必须改成 7否则前向传播根本跑不通。通常还需要把输入尺寸同步调整ViT 是 224x224如果数据集里的图分辨率太低直接把 ViT 的 Patch Embedding 层换小也是常规操作。5. 避坑指南小样本训练中四个常见翻车点与排查思路训练集和测试集划分好了不代表数据就是干净的。以下是我在这个数据集上实际踩过或者见过的四个坑每一条都按“现象 → 原因 → 解决”来记录遇到的可以直接照方抓药。5.1 坑 1训练集和测试集里同一类图片分辨率不一致现象训练 loss 下降得很顺但验证集准确率卡在 60% 左右上不去。原因数据集中部分图片是 224x224部分来自手机拍摄是 512x512。经过Resize(256)和CenterCrop(224)之后不同分辨率图片里的物体占画布比例完全不一样。分辨率高的图缺陷区域占据的画面更大分辨率低的图缺陷可能被裁掉一半特征分布出现偏差。解决统一数据质量写脚本检查所有图片的尺寸把低于某个阈值比如 200x200的图片单独挑出来看。如果缺陷明显做超分重建后再进训练集如果不清晰干脆丢弃不要让它污染模型学习。5.2 坑 2Dry Rot 和 Pink Rot 视觉特征相似模型收敛慢现象训练了 50 个 epoch整体准确率尚可但混淆矩阵里 Dry Rot 和 Pink Rot 两类互相错判的概率明显高于其他类别。原因这两类缺陷在颜色上都是深褐色纹理上都是干枯区域类间距离非常近。交叉熵损失在优化时把它们当作完全独立的类别处理模型很难学到足够区分的细节特征。解决换用 Focal Loss 或 Label Smoothing。Focal Loss 会让模型把学习重心放到难分样本上Label Smoothing 则避免模型对训练集类别过度自信给错误预测留一点缓冲空间。在实际项目中我通常会用LabelSmooth简单改一下损失函数的计算方式。5.3 坑 3JSON 文件里类别顺序与文件夹顺序不一致现象训练时 loss 正常下降但打印出来的预测标签语义完全错乱比如预测结果和真实结果对不上准确率却还行。原因在构建Dataset时如果用os.listdir遍历文件夹得到类别列表再映射到 JSON 里的类别名两者的顺序很可能不一致。比如文件系统里第一个文件夹是Blackleg而 JSON 里第一个类别是Dry Rot索引直接错位。解决只在代码里维护一张唯一的类别映射表从 JSON 里读取类别名按字母排序后生成class_to_idx训练和推理都用这份映射不要用文件系统的遍历顺序。建议把它保存成classes.txt每行一个类别名保证后续推理阶段不会混乱。import json # 读取类别并建立唯一索引映射 with open(label.json, r, encodingutf-8) as f: labels json.load(f) classes sorted(set(labels.values())) class_to_idx {c: i for i, c in enumerate(classes)} # 将映射关系落盘推理阶段直接读取双重保险 with open(classes.txt, w) as f: for c in classes: f.write(c \n)5.4 坑 4小样本下验证集波动大单次准确率不可信现象同一套代码和参数第一次跑验证集准确率 92%第二次随机种子一变就只有 80%。原因整个测试集可能只有 50 到 80 张图其中某一张比较刁钻的图被预测错了准确率就会掉 1 到 2 个点。单次实验的准确率波动范围大完全不能说明模型优劣。解决固定所有随机种子PyTorch、NumPy、Python 内置的random然后用 5 次重复实验取平均值。如果追求更严谨直接在这个数据集上做 K 折验证把小样本的评估误差摊平。6. 进阶玩法1-shot 评估数据集与 TorchScript 模型导出450 张数据集除了做标准分类训练其实非常适合做小样本分类评估。很多刚接触的小白会问1-shot 和 5-shot 到底怎么测我把常用的做法拆给你看。6.1 用 1-shot / 5-shot 评估数据集的表征能力小样本图像分类的核心思路是模型不是直接预测类别而是学习把图片映射成特征向量。测试时每个类别只拿出 1 张或 5 张图作为“支持集”计算它们的特征中心然后将待测图片的特征与这些中心做距离比较离谁最近就归为哪一类。这里直接复用了训练好的 backbone 来抽取特征。import torch import torch.nn.functional as F torch.no_grad() def extract_features(model, loader, device): model.eval() feats, labels [], [] for images, targets in loader: images images.to(device) feat model.backbone(images) # 取骨干输出 (B, 2048, 1, 1) feat model.pool(feat).flatten(1) # 拉平成 2048 维向量 feats.append(feat.cpu()) labels.append(targets) return torch.cat(feats), torch.cat(labels) # 假设 features 形状为 (N, 2048), labels 形状为 (N,) for shot in [1, 5]: support_idx, query_idx [], [] for c in range(7): c_idx (labels c).nonzero().squeeze(1) support_idx.append(c_idx[:shot]) # 每个类别取前 shot 张 query_idx.append(c_idx[shot:]) # 其余做查询 support_feat features[torch.cat(support_idx)] query_feat features[torch.cat(query_idx)] query_label labels[torch.cat(query_idx)] # 计算查询集与支持集特征之间的欧氏距离 dist torch.cdist(query_feat[:, None, :], support_feat[None, :, :]).squeeze(1) pred dist.argmin(dim1) acc (pred query_label).float().mean().item() print(f{shot}-shot 准确率: {acc:.3f})torch.cdist直接暴力计算两两距离数据量小所以速度不慢。model.backbone和model.pool来自第 4 章的ImprovedResNet50这就是为什么我当时强调要把池化层显式暴露出来——抽取特征时会非常方便不用再走一个完整的forward。6.2 导出 TorchScript 模型脱离 Python 环境的推理方案如果你训练好的模型需要交给别人使用而对方环境里没有 PyTorchTorchScript 就是最稳妥的导出方案。把模型trace成静态图对方只需要用 LibTorch 就能直接加载推理。# 导出前必须切换成 eval 模式固定 BN 和 Dropout 的行为 model.eval() dummy_input torch.randn(1, 3, 224, 224) traced_model torch.jit.trace(model, dummy_input) traced_model.save(potato_classifier.pt) print(模型已导出为 potato_classifier.pt)这里用的是torch.jit.trace而不是torch.jit.script。原因是script需要保证模型代码完全被 TorchScript 编译器兼容而我们的 CBAM 模块里有 Python 的list拼接操作script容易翻车。trace只会记录实际执行路径把经过的算子固化下来对含自定义模块的场景友好得多。从那以后我每次拿到一个新的分类数据集都会强制走一遍固定流程先校验 JSON 标签和文件夹对应关系再跑一次 ResNet50 基线确定下限最后才轮到各种改进模块上场——顺序反了改进效果到底是模型带来的还是数据波动带来的根本说不清。这个数据集正好能帮你把这条链路完整跑通希望帮到你。本文还有配套的精品资源点击获取
