简介芒果成熟度图像分类数据集是一份面向图像分类任务的已标注数据包包含约9000张芒果图片划分为成熟、未成熟、损坏三个类别适配果蔬成熟度识别、农产品质检等场景。压缩包内共2000个文件以1998个jpg图像为主体各类别图片按目录存放并预先划分出训练集、验证集和测试集另有1个json分类标签文件、1个show.py可视化脚本整体约261.44MB。目前已有62人学习下载。json文件提供了类别与标签的映射关系show脚本可随机展示样本帮助快速核查数据质量从目录结构到标签配置都经过整理研究者可直接用于图像分类模型训练、对比实验或教学演示开箱即用。无论是训练轻量级卷积网络还是验证Transformer分类模型都能节省大量数据准备时间。1. 9,000张芒果成熟度标注图值不值得投入生鲜分拣线上芒果成熟度的判断一直是最难标准化的环节。老师傅靠捏、摸、看新人靠跟线学换一条生产线、换一个光源标准就漂移。我拿到这个「芒果成熟度图像分类数据集【已标注约9,000张数据】」时第一反应不是马上训模型而是先算一笔账9,000张已标注图对标当前图像分类技术足够训练出可用的分类器吗答案是够。但要看你如何划分训练集、如何处理类别不均衡以及怎么避免模型记住芒果旁边的纸箱而不是芒果本身。这篇文章我会从数据集拆解开始用PyTorch把一条完整可复现的落地链路讲清楚包括参数怎么设、坑在哪、验证怎么做。适合做农业视觉、智能分拣、工业质检的工程师也适合想用图像分类做毕设但不想被数据集糊弄过去的人。2. 芒果成熟度数据集怎么用标签划分与DataLoader实现2.1 成熟度标签怎么定先搞清楚类别含义和数量分布这个数据集标注了约9,000张图像但「成熟度」到底分几类决定了后面所有设计。常见做法是三类未熟、半熟、成熟也有做到四类的比如加一个过熟。拿到数据集第一件事不是直接开训而是统计每个类别的样本数。我一般会写一个十行脚本扫一遍目录。from collections import Counter import os data_dir mango_dataset counts Counter() for sub_dir in os.listdir(data_dir): sub_path os.path.join(data_dir, sub_dir) if not os.path.isdir(sub_path): continue counts[sub_dir] len(os.listdir(sub_path)) print(counts)这个脚本会输出类似Counter({unripe: 3200, half_ripe: 2800, ripe: 3000})的分布。注意如果输出里某一类比其余少一半以上后面训练要立刻处理类别加权否则模型会无脑预测多数类。2.2 从图像目录到PyTorch DataLoader一套通用的加载方案数据集的标注形式可能是按文件夹组织的也可能是CSV或JSON文件。我习惯先把所有标注归并成一张表再统一构造Dataset。以最常见的文件夹结构为例每张图的标签就是它的父目录名。import torch from torch.utils.data import Dataset from PIL import Image class MangoDataset(Dataset): def __init__(self, data_dir, transformNone): self.samples [] self.classes sorted(os.listdir(data_dir)) self.class_to_idx {cls: i for i, cls in enumerate(self.classes)} for cls in self.classes: cls_dir os.path.join(data_dir, cls) for fname in os.listdir(cls_dir): if fname.lower().endswith((.jpg, .jpeg, .png)): self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls])) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] image Image.open(path).convert(RGB) if self.transform: image self.transform(image) return image, label这里有个细节convert(RGB)不能省。很多农副产品图像是RGBA或灰度图直接把图片塞进预训练模型会让第一层卷积报形状错误。PyTorch的DataLoader配合这个Dataset时建议设num_workers4以上否则9000张图的读取会成为训练瓶颈。如果数据集的标注是CSV只需要把__init__里读文件的部分改成pandas.read_csv用文件名和标签列构造self.samples即可。要注意CSV里的文件名可能带相对路径前缀统一用os.path.join拼接别在Windows和Linux上踩斜杠的坑。2.3 训练/验证/测试集划分分层抽样与随机种子很多人拿到数据集后直接用random.sample切分我劝你别这么做。芒果成熟度数据通常存在拍摄批次差异早批的芒果和晚批的芒果在颜色分布上天然不同这会让随机划分把同一批次的图像同时分进训练和验证验证分数虚高。正确做法是分层抽样让每种类别在训练、验证、测试集中的比例保持一致。from sklearn.model_selection import train_test_split X [s[0] for s in dataset.samples] y [s[1] for s in dataset.samples] train_paths, test_paths, train_labels, test_labels train_test_split( X, y, test_size0.2, stratifyy, random_state42) train_paths, val_paths, train_labels, val_labels train_test_split( train_paths, train_labels, test_size0.2, stratifytrain_labels, random_state42)stratifyy是关键参数它保证每个类别在切分后仍保持原始比例。random_state42是为了复现。如果最终训练集与验证集出现光晕、阴影差异过大说明批次效应已经被带进来了后面需要做进一步的颜色归一化。这里多花十分钟能避免训练到一半才发现验证集和训练集长得不像同一个芒果园。3. 跑通第一个分类模型ResNet迁移学习训练全流程3.1 迁移学习选型为什么先选ResNet18而不是DeiT9,000张图像做图像分类量级不大不小。现在最新的图像分类模型层出不穷Swin Transformer、ConvNeXt、EfficientNet都很能打但我不建议一上来就跑它们。原因很朴素你是来落地成熟度分类的不是来刷基准的。ResNet18参数量约1100万在ImageNet上预训练权重容易获取训练一个epoch用不了几分钟而且足够拟合9000张数据。如果你非要用EfficientNet我建议用B3以上的版本B0太轻解决不了成熟度这种细粒度分类问题ViT类模型在没有做足够数据增强的情况下翻车概率远高于CNN。3.2 数据预处理与增强让模型不再「只认这张芒果」分拣线的光源、相机角度、输送带速度都会让图像分布变化。所以数据增强不是锦上添花而是防止过拟合的刚需。我常用的transforms如下import torchvision.transforms as T train_transform T.Compose([ T.RandomResizedCrop(224, scale(0.7, 1.0)), T.RandomHorizontalFlip(p0.5), T.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.05), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform T.Compose([ T.Resize(256), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])RandomResizedCrop(224, scale(0.7, 1.0))会随机缩小芒果占画面的比例迫使模型关注果实本身而不是固定背景。ColorJitter的hue参数别调太大芒果成熟度反映在颜色从绿到黄褐的过渡上色相抖动太过会抹掉成熟度的区分度。这里最容易犯的错是验证集也用了训练集的增强。验证集必须只用确定性变换否则每次验证打的分数都是随机抖动你根本看不出模型是进步还是退步。3.3 训练循环核心代码从加载权重到学习率调度训练脚本我习惯这样组织先加载预训练ResNet18并替换最后一层全连接然后定义损失函数、优化器和学习率调度器。import torch import torch.nn as nn from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20)lr1e-4是迁移学习的默认安全值。如果你只替换全连接层前面层冻结学习率可以放到3e-4全模型微调时降回1e-4。CosineAnnealingLR会把学习率从初值余弦降到接近0比我以前用的StepLR稳定得多尤其是在训练后期不会因为学习率骤降而震荡。训练循环我会单独写一个函数支持混合精度和梯度裁剪from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for epoch in range(20): model.train() total_loss 0.0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() with autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() total_loss loss.item() * images.size(0) print(fepoch {epoch1}, loss {total_loss/len(train_loader.dataset):.4f})clip_grad_norm_在迁移学习里经常被省略但在小数据集上非常管用。当某个batch里面的芒果图像恰好全是极端光照梯度会突然爆掉裁剪到1.0能保住整个已经训好的主干特征。训练集9000张图batch size设3220个epoch在单张消费级显卡上大约30分钟跑完。跑通这一步你手里就有了一个能看的基线。4. 从90%到95%的调优损失函数、解冻策略与交叉验证4.1 类别不平衡换用带权重的损失函数基线跑通后验证集准确率大概率在88%~93%之间前提是类别分布没有太离谱。如果Counter显示某一类只占10%那损失函数必须换。最直接的做法是给交叉熵加类权重from sklearn.utils.class_weight import compute_class_weight import numpy as np class_weights compute_class_weight(balanced, classesnp.unique(y_train), yy_train) class_weights torch.tensor(class_weights, dtypetorch.float).cuda() criterion nn.CrossEntropyLoss(weightclass_weights)compute_class_weight会根据样本数自动计算权重样本少的类别获得更高的惩罚。这是最简单的方案。如果你发现加上权重后多数类的准确率掉得太多说明类别之间其实存在可混淆性比如「半熟」和「成熟」边界本来就模糊。这时可以考虑用label_smoothing0.1让模型不要对某个类别过于自信同时缓解权重带来的震荡。4.2 微调节奏先冻结主干再解冻低层迁移学习最大的玄学在于「什么时候解冻哪一层」。很多人一上来就全参数微调结果训练集准确率不到5个epoch就冲到98%验证集却一直趴在90%。原因在于主干已经被打乱低层特征在重新适应小数据集时产生了过拟合。我一般分两段走第一阶段冻结除model.fc外的所有层只训练新全连接层。因为新分类头的梯度不会传回主干预训练特征不会被破坏。for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True optimizer torch.optim.AdamW(model.fc.parameters(), lr3e-4) # 训练5个epoch左右让分类头先贴合芒果特征第二阶段解冻主干中最后的Block4和fc其余继续冻结。注意把BatchNorm层设为训练模式否则running_mean不更新解冻效果大打折扣。for param in model.layer4.parameters(): param.requires_grad True for module in model.modules(): if isinstance(module, nn.BatchNorm2d): module.training True optimizer torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr5e-5)这里的学习率要比第一阶段低得多因为主干特征已经稳定只需要微调让它更适应芒果的纹理和颜色。注意model.eval()会在验证阶段自动关闭BN更新但如果在训练阶段用了torch.no_grad()或者自定义推理循环很容易把BN的training状态弄错导致loss曲线正常、验证acc却一直不升高。4.3 K折交叉验证用9000张数据榨出更稳的模型如果你要在分拣线上正式投入使用单次划分的验证准确率不足以让你下决心。我建议在调参结束后用5折交叉验证复训一遍这样能拿到每个类别的平均准确率和方差而不只是撞运气的那一个分数。from sklearn.model_selection import StratifiedKFold skf StratifiedKFold(n_splits5, shuffleTrue, random_state42) fold_scores [] for fold, (train_idx, val_idx) in enumerate(skf.split(X, y)): train_subset torch.utils.data.Subset(dataset, train_idx) val_subset torch.utils.data.Subset(dataset, val_idx) # 重新构造DataLoader并训练 fold_score run_training(train_subset, val_subset) fold_scores.append(fold_score) print(f5折均值: {np.mean(fold_scores):.2f}% 标准差: {np.std(fold_scores):.2f}%)StratifiedKFold和train_test_split的区别在于它会保证每一折的类别比例一致但批次效应可能让某一折特别难。如果某折的分数比其余低5个点以上那大概率是这折里面包含了一批光照偏暗的图像。别急着放弃这个fold把它当成需要在增强里加入T.Grayscale(3)或T.RandomAdjustSharpness的信号。5. 训练避坑标注噪声、过拟合与分类翻车的4个排查记录5.1 现象训练loss下降验证acc每5个epoch就掉一次我遇到过最典型的情况是训练loss一路下滑到0.1验证acc在92%附近上下抖动但每隔4~5个epoch突然跌到80%。最终排查发现学习率调度用了StepLR(step_size4, gamma0.5)每个step之后学习率半衰验证集准确率跟着跳一下。原因余弦退火虽然有波动但不会突变StepLR阶跃式降低学习率模型在边界处参数跳变如果此时恰好跑到一个偏难batch验证分数就会被拉低。解决方法是换成CosineAnnealingLR(T_max20)或者把gamma改成0.9并让step_size覆盖完整epoch让学习率平缓下降。5.2 现象模型对「半熟」这个类别永远预测不准按成熟的颜色卡半熟处于绿转黄的过渡段视觉特征和两侧类别重叠度很高。我刚开始训的时候半熟类的召回率只有65%其余全部分到了未熟或成熟。原因有两层一是半熟的真实标注本身就带噪声不同标注员对「半熟」判断的边界不统一二是类别本身存在模糊性属于标签噪声和类间重叠叠加。解决思路是接受「半熟」的不可分性把损失函数从CrossEntropyLoss换成FocalLoss(gamma2.0)降低易分类样本对loss的主导让模型更关注过渡区的样本。class FocalLoss(nn.Module): def __init__(self, gamma2.0, alphaNone): super().__init__() self.gamma gamma self.alpha alpha def forward(self, logits, targets): ce_loss nn.functional.cross_entropy(logits, targets, reductionnone) pt torch.exp(-ce_loss) focal_loss (1 - pt) ** self.gamma * ce_loss if self.alpha is not None: focal_loss focal_loss * self.alpha.gather(0, targets) return focal_loss.mean()gamma2.0是图像分类里最常用的默认值。alpha可以代入选好的类别权重。换掉之后半熟类的召回率通常能提升10个点以上但要注意多数类的准确率会略微下降这是合理的代价。5.3 现象验证集准确率很高但换一台手机拍芒果就崩在实验室数据上跑了94%准确率拿去分拣线现场一试大量误判。我后来把测试图像调出来看发现模型关注的是图像右下角的纸箱纹理而不是芒果皮色。原因出在数据采集时背景单一训练集所有图片都在同一个台面上拍的模型把背景当成了判别特征。解决方法是回炉数据筛查把每张训练图都做一次小范围随机裁剪让芒果边缘和背景频繁出现截断同时增加HorizontalFlip。另外我养成了一个习惯每轮训练完取错误案例画一个两列的对比图左边是错图右边是对应的预测概率条。如果错图里全是同一个背景基本就是背景过拟合。5.4 现象训练集准确率99%验证集却只有88%这是最典型的过拟合信号而且多发生在没有用预训练权重或者学习率设成1e-3的情况。9,000张图像不算少但对于ResNet18来说仍然不足以从头学习通用特征。原因模型容量过大训练集信息被完整记忆。解决路径有三步先加载ImageNet预训练权重再把weight_decay从1e-4提到5e-4最后把RandomResizedCrop的scale下限从0.7降为0.4迫使模型看到更多局部细节。做完这三步验证集acc通常能回到92%上下且不会出现训练集acc远高于验证集的情况。6. 部署前的验证手法混淆矩阵、Grad-CAM与ONNX导出验证成熟度分类模型我不只看总体准确率而是用三件套混淆矩阵、Grad-CAM、ONNX导出。先算混淆矩阵把每个类别的真实与预测结果做成矩阵。如果「半熟」和「成熟」之间有大面积互混说明分拣线需要增加物理约束比如传送带速度配合多角度拍摄而不是继续调模型。from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns preds, truths [], [] model.eval() for images, labels in val_loader: images images.cuda() pred model(images).argmax(dim1).cpu().numpy() preds.extend(pred) truths.extend(labels.numpy()) print(classification_report(truths, preds, target_namesclass_names)) sns.heatmap(confusion_matrix(truths, preds), annotTrue, fmtd)Grad-CAM可以快速定位模型看的是芒果果皮还是叶片阴影。用pytorch_grad_cam库时我习惯选择layer4[-1]和layer4[-2]两个层分别可视化。如果热力图集中在芒果边缘而不覆盖大块果皮说明模型仍在依赖轮廓特征成熟度判断会受到果实形状干扰。这个检查做一次比调三天参数都值。部署时导出ONNX是稳妥之举。PyTorch的torch.jit.trace在张量维度变化时容易踩坑ONNX则稳定得多dummy torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy, mango_ripeness.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version17)导出的ONNX文件可以直接交给TensorRT或ONNX Runtime在Jetson这类边缘设备上跑实时推理。我踩过最大的坑是导出前忘记切到model.eval()模式导致导出的模型把BN层也序列化了推理结果与训练验证分数完全不同。所以导出前一定先跑一次验证集确认eval模式下的准确率和推理对齐再执行导出。做完这一步这个数据集才算真正被榨干。我个人的习惯是每次得到一个分类结果会顺手把错误样本输出到error_samples目录下一轮训练前再看一眼。9,000张图说多不多说少不少能帮你快速验证方案但不能帮你解决所有现场噪声。我见过太多项目死在「验证集92%」的自我满足里希望这篇笔记能让你少走几步弯路踏踏实实把模型送到产线上再谈准确率。希望帮到你。本文还有配套的精品资源点击获取