ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

小样本图像分类实战:用ResNet18迁移学习识别马铃薯缺陷

小样本图像分类实战:用ResNet18迁移学习识别马铃薯缺陷 简介面向图像分类入门、农产品缺陷检测与模型对比实验等场景这份马铃薯缺陷标注数据集适合初学者快速上手分类模型也便于研究者做算法对比与调优。数据涵盖Dry Rot、Blackleg、Pink Rot等7类缺陷共约450张图像已按类别划分训练集与测试集各目录均存放同一类样本并附json标注文件与show可视化脚本可直观查看每一类的样本分布方便直接训练、验证与调试。压缩包内共454个文件以jpg/jpeg图像为主体另有1个json标注文件、1个Python可视化脚本和1张说明图片整个资源包仅12.58MB轻量灵活适合快速实验。目前已有150人学习配套专栏还介绍了CNN分类网络的改进方法可在该数据集基础上进一步探索结构优化与精度提升是练习图像分类流程的实用小成本选择整体方案完整解压后即可按类别浏览训练与测试样本。1. 马铃薯食物缺陷图像分类数据集450张标注图能训练出什么马铃薯食物缺陷图像分类数据集450 张已标注图像在图像分类数据集里算微缩但足够验证一条完整的落地流程数据划分、类别统计、迁移学习、缺陷识别、结果可视化。它解决的不是刷榜而是用最小成本验证食品质检路线的可行性。这套数据常见按健康、褐变、绿皮、发芽、腐烂打单标签适合刚入门的工程师走通第一个真实图像分类项目也适合团队在凑不齐上万张样本时做方案预演。结论先行数据划分干净、增强克制的前提下ResNet18 迁移学习可以把五类缺陷稳定做到 90% 上下。下面从拿到数据后的第一件事开始讲。2. 先做数据集体检再谈训练目录结构、标注格式与类别分布2.1 目录结构与标注格式先摸清数据是怎么躺着的拿到一个图像分类数据集第一步不是写模型而是把数据目录完整列一遍。常见做法是 ImageFolder 风格——每个类别一个文件夹文件夹名就是标签。马铃薯缺陷数据集通常长这样potato_defect/ ├── train/ │ ├── healthy/ # 健康薯 │ ├── bruise/ # 磕碰褐变 │ ├── green/ # 绿皮 │ ├── rot/ # 腐烂 │ └── sprout/ # 发芽 ├── val/ └── test/如果拿到的是 CSV 标注一般就是 filename,label 两列配合一个存图片的目录。这种格式也不难转按 label 建文件夹把图片拷进去就行。标注文件长什么样不重要重要的是统一成 ImageFolder 结构之后PyTorch 的datasets.ImageFolder可以直接读取省掉手写 Dataset 的功夫。图像格式以 JPG 为主分辨率常见在几百到一千像素之间。450 张图意味着平均每类 90 张如果某一类只有二三十张后面类别不均衡的处理就必须做。先看目录结构能快速判断数据是按类别分好的还是散落一地需要自己整理。空目录也要顺手找出来ImageFolder遇到空目录会直接报错而且错误信息并不直观。2.2 类别分布统计脚本训练前先做一次体检统计类别分布是小样本图像分类最容易被跳过、又最该先做的一步。下面这段脚本遍历训练目录输出每个类别的样本数和占比from pathlib import Path from collections import Counter data_root Path(potato_defect/train) if not data_root.exists(): raise FileNotFoundError(训练目录不存在先确认解压路径) counts Counter() for cls_dir in data_root.iterdir(): if cls_dir.is_dir(): counts[cls_dir.name] len(list(cls_dir.glob(*.jpg))) total sum(counts.values()) print(f总样本数: {total}) for cls_name, n in counts.most_common(): print(f{cls_name}: {n} 张, 占比 {n / total * 100:.1f}%)这段脚本做了什么遍历 train 下每个子目录用glob统计以 .jpg 结尾的文件数。Counter按数量从大到小排序输出每类样本量与占比。逻辑很简单但它决定后面所有决策——哪类样本少、要不要加权采样、增强重点照顾谁全看这个输出。几个容易漏的细节。glob(*.jpg)只匹配 .jpg如果数据里混了 .png 或 .jpeg统计会偏小建议把.jpg、.jpeg、.png分开统计后相加。如果某目录里混了 Thumbs.db 之类的隐藏文件用iterdir()数会出错glob不受影响。还有统计完顺便确认没有空目录——空目录在ImageFolder读取时会报错而且错误信息指向不明确。2.3 单标签还是多标签一个马铃薯可能同时有几种毛病真实场景里缺陷不会按教科书长。一个薯可能既发芽又在局部腐烂绿皮的薯同时带着磕碰。450 张的标注数据大部分会按最明显、最影响售卖的缺陷打单标签。这是最省事的做法单标签直接用nn.CrossEntropyLoss模型输出层不用改结构。但单标签有代价。同一张图存在两种缺陷时标签只反映其中一种模型训练时会看到标着发芽的图里有大块腐烂区域梯度信号被污染。我一般拿到数据后会抽查每类图片把明显包含多种缺陷的样本挑出来要么删掉要么归到更严重的类别。多标签方案不是不行但要换成nn.BCEWithLogitsLoss输出层从 softmax 改成 sigmoid验证指标也要从 accuracy 换成 F1 或 mAP。450 张图做多标签训练样本量偏紧除非每个缺陷类别都有足够多的正样本否则不建议一上来就上多标签。先把单标签跑通再考虑要不要扩展。3. 用迁移学习跑通马铃薯缺陷分类ResNet18 微调最小流程与参数3.1 为什么不从零训练450张图扛不住随机初始化450 张图从随机初始化训练 ResNet18结果几乎可以预见过拟合。ResNet18 有大约 1100 万参数450 张图连参数数量的零头都不够。从零训练时底层卷积要重新学边缘、纹理、颜色块这些基础特征这点数据量远远不够。迁移学习的逻辑是用 ImageNet 预训练权重做起点这些权重已经把通用视觉特征学好了。马铃薯缺陷——褐变、绿皮、腐烂——本质上是颜色和纹理的局部变化预训练模型底层的边缘检测器、纹理检测器可以直接复用。真正要训练的是最后的分类层让它把通用特征组合映射到健康/褐变/绿皮/发芽/腐烂五类上。迁移学习不是无脑全量微调。小数据集上全量微调风险很大容易把预训练权重冲坏。我一般分两步第一步冻结 backbone只训练新增的 fc 层第二步解冻最后两个 stage用小学习率微调。450 张图通常第一步就能达到不错效果第二步只做锦上添花。3.2 最小可运行的训练脚本从加载到保存代码从数据加载到模型构建一次给齐注释标了关键点import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms, models device torch.device(cuda if torch.cuda.is_available() else cpu) train_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomRotation(15), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_ds datasets.ImageFolder(potato_defect/train, transformtrain_tf) val_ds datasets.ImageFolder(potato_defect/val, transformval_tf) train_loader DataLoader(train_ds, batch_size16, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size16, shuffleFalse, num_workers4, pin_memoryTrue) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, len(train_ds.classes)) model.to(device)核心是ImageFolder直接读目录结构标签由文件夹名自动生成类名按字母顺序排序。所有输入统一缩放到 224×224这是 ImageNet 预训练模型的标准输入尺寸。Normalize用的均值方差也是 ImageNet 统计值不能按自己数据集重算否则预训练权重就对不上了。提示新版本的 torchvision 推荐用weights...而不是pretrainedTrue前者会明确指定权重来源后者在较新版本中会告警并可能被移除。训练循环criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3) def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, correct, total 0.0, 0, 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, preds outputs.max(dim1) correct (preds labels).sum().item() total labels.size(0) return total_loss / total, correct / total for epoch in range(20): loss, acc train_one_epoch( model, train_loader, optimizer, criterion, device) print(fepoch {epoch1:02d} | loss {loss:.4f} | acc {acc:.4f}) torch.save(model.state_dict(), fpotato_epoch_{epoch1:02d}.pth)注意optimizer.zero_grad()必须在每次前向之前调用否则梯度会累加导致 loss 忽大忽小。outputs.max(dim1)取每个样本得分最高的类别索引作为预测。每个 epoch 保存一次权重是为了做早停或回滚——验证集最优的 epoch 未必是最后一个多存几个 checkpoint 相当于给自己留一颗后悔药。3.3 训练参数参考表学习率、batch size和轮数参数参考值说明batch size16450 张图一个 epoch 约 28 步BN 统计够用初始学习率1e-3只训FC / 1e-4微调backbone学习率大容易震荡epoch20 起步看验证集准确率是否还在涨优化器Adam 或 SGD momentum(0.9)Adam 收敛快SGD 更稳输入尺寸224×224与预训练输入一致别随意改大学习率是最敏感的参数。只训练 fc 层时 1e-3 通常没问题一旦解冻 backbone学习率必须降到 1e-4 量级否则几个 epoch 就把预训练权重冲坏了。batch size 也别贪大450 张图用 64 的 batch一个 epoch 只有 7 步BatchNorm 的统计量会非常不稳验证集指标跟着乱跳。4. 数据增强与类别不均衡450张小样本分类的四个关键参数4.1 数据增强的幅度旋转翻转随便用颜色抖动要克制数据增强对 450 张图几乎是救命稻草。随机旋转和水平翻转这类几何变换对马铃薯缺陷数据集相对安全——旋转 15 度不会改变这是不是绿皮的事实。但颜色类增强要非常克制这是我反复翻车后总结出来的。马铃薯的很多缺陷本身就是颜色问题绿皮是表皮叶绿素累积褐变是酚类氧化腐烂是灰褐色软烂区域。如果用ColorJitter把色调、饱和度大幅随机抖动可能把一张绿皮薯的颜色抖成正常薯标签就变成错的了。小样本图像分类里这种错误标签对训练的破坏比欠拟合还严重。建议的增强组合是随机旋转 15 度以内、水平翻转、随机缩放 0.9 到 1.1 倍。颜色域只做很小的亮度和对比度扰动不动 hue。用 torchvision 写train_tf transforms.Compose([ transforms.RandomResizedCrop((224, 224), scale(0.9, 1.1)), transforms.RandomRotation(15), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.1, hue0.0), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ])亮度 0.2、对比度 0.2 的扰动是为了模拟不同拍摄环境的光线差异这是真实质检场景必然存在的。饱和度只给 0.1 的小幅度hue 直接设 0。从我跑过的经验看hue 一旦超过 0.05绿皮类和健康类的混淆率会明显上升因为模型能用来区分两者的颜色线索被增强人为抹掉了。4.2 类别不均衡的两种解法加权采样还是改损失函数450 张图里健康薯占 60%、腐烂薯只有 15% 是很常见的情况。直接训CrossEntropyLoss腐烂类会学得很差因为多数类的梯度贡献大少数类的错误对总 loss 影响太小。第一种解法是加权采样用WeightedRandomSampler让少数类在每轮被多抽几遍from torch.utils.data import WeightedRandomSampler import torch labels [lb for _, lb in train_ds.samples] class_counts torch.bincount(torch.tensor(labels)) sample_weights 1.0 / class_counts[labels].float() sampler WeightedRandomSampler(sample_weights, num_sampleslen(labels), replacementTrue) train_loader DataLoader(train_ds, batch_size16, samplersampler)关键在sample_weights每个样本的权重等于它所在类别样本数的倒数。假设健康薯 270 张、腐烂薯 45 张腐烂薯每个样本的采样权重是健康薯的 6 倍。replacementTrue表示有放回采样少数类的样本在一轮里会重复出现。第二种解法是给损失函数加类别权重改动更小class_weights 1.0 / class_counts.float() class_weights class_weights / class_weights.sum() * len(class_counts) criterion nn.CrossEntropyLoss(weightclass_weights.to(device))两种方案可以共存但我倾向先用损失权重因为改动小、好调试不影响数据加载逻辑。如果加了权重少数类还是学不动再叠加加权采样。要注意CrossEntropyLoss的weight参数必须在构造时就放到和模型相同的设备上否则 CPU/GPU 张量不匹配会直接报错。4.3 四个必调参数的参考值参数推荐范围作用旋转角度10°~20°模拟薯在传送带上的姿态变化色相扰动 hue0绿皮/褐变是颜色缺陷不能动类别权重按样本数反比解决不均衡防多数类通吃FC层 dropout0.3~0.5配合全量微调小样本全量微调时的过拟合防线第四个参数 dropout 值得单独说。默认 ResNet 的 fc 层不带 dropout小样本全量微调时最后几层很容易过拟合。我一般把 fc 改成nn.Sequential(nn.Dropout(0.4), nn.Linear(in_features, num_classes))。如果只训练 fc 层dropout 意义不大一旦解冻 backbone 做微调这个 0.4 的 dropout 能明显稳住验证集。5. 小数据集图像分类避坑实录5条踩坑记录5.1 训练集验证集高准确率测试集崩盘现象训练集和验证集都跑到了 95% 以上拿现场新拍的照片一测准确率掉到 70% 出头。原因数据划分没按拍摄来源分。同一个马铃薯从不同角度拍多张图随机划分会让同一颗薯的照片同时出现在训练和验证集里。模型记住的是这颗薯的样子而不是这类缺陷的共性验证集高分是假象。解决按拍摄批次或薯个体分组划分。具体做法是先给文件名加个体前缀比如potato_001_angle1.jpg、potato_001_angle2.jpg然后按前缀分组一组整体划入训练集或验证集绝不允许同一个前缀横跨两个集合。5.2 训练loss稳定下降验证loss来回震荡现象前 5 个 epoch 训练 loss 降得很顺验证 loss 从第 3 个 epoch 开始上下乱跳准确率跟着抖。原因学习率偏大模型在最优解附近横跳另一种常见情况是只训练 fc 层backbone 里的 BatchNorm 还在用预训练统计量验证阶段输入分布一偏移就出错。解决先挂ReduceLROnPlateau让 loss 不降时自动降学习率factor 设 0.5。如果震荡依旧解冻layer3、layer4两个 stage学习率降到 1e-4并且把模型设为 train 模式重新统计 BN。5.3 模型把一切都预测成多数类现象健康薯占 60% 以上时训练出来的模型对任何输入都输出 healthy。验证集准确率停在 60% 上下看似不低实则完全不能用。原因交叉熵损失在类别不均衡时天然偏向多数类少数类样本虽然被正确分类的收益高但错误分类的惩罚太小梯度被多数类淹没。解决给CrossEntropyLoss传weight参数按样本数反比设置。加了权重之后腐烂薯的预测错误会带来更大的梯度模型才会真正去学少数类特征。做完这一步再看分类报告里每个类别的 F1而不是只看准确率。5.4 加了颜色增强后准确率反而下降现象在增强里加入ColorJitter之后训练集准确率下降验证集也降比不加增强还差。原因色调抖动范围太大。绿皮薯的颜色被抖到了正常薯的色域里等于人为制造了错误标签。数据增强不是越多越好它必须不改变语义。解决把 hue 设为 0saturation 控制在 0.1 以内只保留小幅度亮度和对比度扰动。这条原则对任何颜色本身就是缺陷特征的数据集都适用马铃薯、辣椒、肉类质检都一样。5.5 450张图跑出验证集100%准确率现象模型训练完验证集准确率 100%怎么看都不真实换了测试集立刻掉到 80% 以下。原因极可能是数据泄漏。拍摄时同一个薯的连续帧、同一袋薯在不同光源下的照片被随机切分到了训练集和验证集。验证集里存在训练集的孪生图模型靠记忆而不是泛化拿到了满分。解决回查验证集图片的文件名和路径确认没有同一来源的重复。更稳妥的做法是划分时按文件名前缀分组连续编号的样本视为同一个体整体落入同一集合。数据划分这件事在小数据集上花半小时做干净比后面调三天参数都值。6. 用混淆矩阵和Grad-CAM验证模型学到了什么缺陷特征训练跑完别急着收工。准确率只是一个数字你得知道模型到底把哪两类搞混了、决策依据又是什么。混淆矩阵和 Grad-CAM 是小样本图像分类最值得做的两个验证动作前者回答错在哪后者回答看的是什么。import numpy as np from sklearn.metrics import confusion_matrix, classification_report import torch model.load_state_dict(torch.load(best_model.pth)) model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) outputs model(images) preds outputs.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_namestrain_ds.classes)) print(Confusion Matrix:\n, cm)看混淆矩阵时先盯两个位置健康薯有没有被大量误报成绿皮腐烂和褐变之间是不是混在一起。前者说明颜色增强还是太激进后者说明纹理特征学得不到位需要解冻更多层补救。classification_report 里每个类别的 F1 值比总体准确率更能反映小类是否被真正学会。Grad-CAM 的作用是可视化模型的关注区域。如果腐烂样本的高亮集中在软烂斑块上说明学到的是缺陷特征如果高亮落在背景阴影、托盘边缘说明模型在用环境线索偷分。这个检查在 450 张图上尤其重要——样本少模型天然容易找捷径。我常用 pytorch-grad-cam 这个库对预测错误的图做可视化每个类别挑三张足够看清模型是不是走正道。我自己的教训是准确率再高都不如把预测错误的几张图用 Grad-CAM 画出来看一眼能省掉大量调参的玄学时间希望帮到你。本文还有配套的精品资源点击获取
返回列表