基于PyTorch的36类果蔬图像分类：数据集解析与迁移学习实战

简介一份面向图像分类入门与实践的常见果蔬多类别标注数据集包含香蕉、苹果、梨、葡萄、橙子、黄瓜、胡萝卜、辣椒、洋葱、土豆等36个类别约4200张已标注图像。数据已完成预处理可直接作为分类网络输入并划分了独立的训练集、测试集与验证集用以支撑模型训练、验证与效果对比省去自行采集、筛选和标注的繁琐步骤。压缩包内共2000个文件以1998张jpg图像为主体附1个Python可视化脚本和1个json类别配置文件整体尺寸约364.87MB。配套脚本可一键展示各类别图片json文件则便于确认类别索引与标签对应关系。这套数据目前已有119人学习使用适合刚接触图像分类的开发者、相关课程实验或小型算法验证项目使用。1. 几千张还是几万张这份 36 类果蔬分类数据集的定位很明确做图像分类的都知道CIFAR-10 太小太玩具ImageNet 又大得跑不动自己爬图光清洗就得耗掉一整天。这份常见果蔬多类别图像分类数据集约 4200 张、36 个类别已经标注、已经划分好训练集/测试集/验证集拿到手就能直接喂给分类网络。香蕉、苹果、梨、葡萄、橙子这些水果加上黄瓜、胡萝卜、辣椒、洋葱、土豆这类蔬菜正好覆盖日常视觉里出现频率最高的果蔬种类。适合三类人做图像分类入门但不想被数据预处理劝退的新手需要小规模基准数据集验证模型改进效果的算法工程师以及做毕业设计需要一份「能跑通还能讲清楚」的数据集的学生。单类平均一百来张不多不少恰好逼着你用迁移学习而不是傻乎乎地从零训练。2. 数据集结构拆解目录规划、json 类别清单与数据划分逻辑2.1 目录划分train/val/test 各自独立同类图片放在同一文件夹这份数据集把「已标注」落实到了文件组织上。训练集、测试集、验证集是三个独立的一级目录每个目录下面按类别建子文件夹同一类别的图片放在同一个子文件夹里。这种组织方式意味着你可以直接用 PyTorch 的torchvision.datasets.ImageFolder或者 TensorFlow 的image_dataset_from_directory加载不用自己写标签映射逻辑。fruit_veg_36/ ├── train/ │ ├── banana/ │ │ ├── Image_78.jpg │ │ ├── Image_84.jpg │ │ └── ... │ ├── apple/ │ │ ├── Image_86.jpg │ │ └── ... │ └── ...共 36 个类别子目录 ├── val/ │ └── ...与 train 相同的类别结构 ├── test/ │ └── ...与 train 相同的类别结构 ├── labels.json # 类别清单 └── show.py # 可视化脚本ImageFolder 加载时会按照文件夹名的字母顺序给类别编号比如apple是第 0 类banana是第 1 类carrot是第 2 类。这里有一个容易踩的坑如果你同时用labels.json里的类别顺序去对应模型输出两个顺序不一致最后画混淆矩阵或者算 per-class accuracy 时全部错位。后面避坑章节会具体说这个问题。2.2 读取 labels.json类别名与文件夹名的映射关系项目描述里明确说了「具体查看 json 文件」所以labels.json就是这份数据集的「类别说明书」。常见做法是 json 里存一个类别名列表或者一个 dict 映射。打开看一眼就知道 36 类到底是哪 36 类。import json with open(labels.json, r, encodingutf-8) as f: labels_data json.load(f) # 常见两种结构实际以资源内文件为准 # 如果是一个列表 if isinstance(labels_data, list): class_names labels_data # 如果是一个 dict比如 {0: banana, 1: apple, ...} elif isinstance(labels_data, dict): # 按 key 排序保证顺序稳定 class_names [labels_data[k] for k in sorted(labels_data.keys(), keyint)] print(f共 {len(class_names)} 个类别) for idx, name in enumerate(class_names): print(f{idx}: {name})注意一点读取 json 后如果 dict 的 key 是字符串形式的数字排序时要用keyint转一下否则10会排在2前面标签顺序全乱。这段代码输出类别列表后我建议先人工核对一遍文件夹名和 json 列表是否一一对应。2.3 类别分布抽样拿到手不要直接开训先数一遍每类有多少张约 4200 张、36 类平均每类大约 117 张。但「平均」不代表「均匀」果蔬这类数据常见的分布不均问题在采集时几乎必然存在——拍摄难度、季节供应、品种多样性都会影响各类别样本量。直接开训的话样本量少的类别会被多数类压制模型对胡萝卜的识别能力可能远低于对香蕉的识别能力。import os from collections import Counter train_root train counts Counter() for class_name in os.listdir(train_root): class_dir os.path.join(train_root, class_name) if os.path.isdir(class_dir): counts[class_name] len([f for f in os.listdir(class_dir) if f.lower().endswith((.jpg, .jpeg, .png))]) # 打印每类数量按从少到多排列 for name, cnt in counts.most_common()[::-1]: print(f{name}: {cnt}) print(f总图片数: {sum(counts.values())}) print(f最少类: {counts.most_common()[-1][0]} ({counts.most_common()[-1][1]} 张)) print(f最多类: {counts.most_common()[0][0]} ({counts.most_common()[0][1]} 张))统计完如果发现最少类和最多类相差超过一倍后面用WeightedRandomSampler或者类别权重来平衡是必要的。这一步几十行代码能省掉后面训练完才发现小类别全被吞掉的返工时间。3. 数据管道搭建从 ImageFolder 到可以直接训练的 PyTorch 输入3.1 用 ImageFolder 读图标签顺序与数据集划分如何对齐目录结构是 train/val/test 各自独立、类别子文件夹平铺这种情况下 PyTorch 的ImageFolder是最优解。它会把train/banana/xxx.jpg这类路径自动映射成(图像张量, 类别索引)对并且类别索引就是按文件夹名的字母序排的。from torchvision import datasets, transforms from torch.utils.data import DataLoader # 训练集做数据增强验证集和测试集只做尺寸调整 transform_train transforms.Compose([ transforms.Resize(256), transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), 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]) ]) transform_eval 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]) ]) train_dataset datasets.ImageFolder(roottrain, transformtransform_train) val_dataset datasets.ImageFolder(rootval, transformtransform_eval) test_dataset datasets.ImageFolder(roottest, transformtransform_eval) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)这段代码里训练集的RandomResizedCrop会随机裁剪出一块区域再缩放到 224×224等效于给模型提供不同缩放、不同构图的样本这在样本量只有一百来张/类的情况下几乎是必需的增强手段。验证集和测试集用CenterCrop保证每次评估的输入是确定的结果可复现。shuffleFalse在验证/测试时也要注意保持否则预测结果和真实标签对不上。3.2 预处理参数为什么是 256 缩放 224 裁剪而不是直接 resize直接transform.Resize((224, 224))会把原本不同宽高比的图片强行拉成正方形苹果变成椭圆香蕉变成直棍形状信息被扭曲模型学到的特征里混入形变噪声。常见的做法是先等比缩放到短边 256再从中心裁一块 224×224这样既保留了主体内容又符合 ImageNet 预训练模型的输入习惯。如果你的数据集里图片尺寸差异很大用Resize(256)后某些图会丢失边缘信息此时可以改用Resize(256, 256)让它统一形变再让网络自己去适应。归一化的 mean/std 直接沿用 ImageNet 的[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]前提是你用了在 ImageNet 上预训练的模型。因为预训练权重的特征提取层是在这套归一化参数下学出来的输入分布保持一致迁移效果才稳定。如果从零训练才需要自己统计数据集的 mean/std。3.3 数据增强配置果蔬分类的增强边界在哪里果蔬分类有一个特殊性颜色是极强区分特征。香蕉的黄色、苹果的红色/绿色、胡萝卜的橙色模型很大程度靠颜色在分类。所以ColorJitter的饱和度、亮度扰动幅度要克制我一般设 brightness0.2、contrast0.2、saturation0.2不会动 hue色相。色相一旦扰动青苹果可能变成红苹果标签却没变等于给模型喂了错误监督信号。翻转增强也要注意方向。左右翻转对果蔬来说基本上是安全的但上下翻转会造成语义异常——草莓上下颠倒虽然还是草莓但光照方向变了模型容易学到「上亮下暗」这种伪特征。训练时只做RandomHorizontalFlip就够了。旋转增强如果要用角度限制在 ±15° 以内比较合理超过 90° 的分类任务会让模型混淆「竖着的香蕉」和「横着的香蕉」。4. 迁移学习训练ResNet18 在这份数据集上的完整参数配置4.1 为什么选 ResNet18 而不是从零训练或直接上 ResNet50约 4200 张图、36 个类别平均每个类一百一二十张。这个规模从零训练一个 ResNet18 很容易过拟合——模型参数远多于样本信息量训练集准确率很快接近 100%验证集却上不去。用 ImageNet 预训练权重做迁移学习是更务实的方案预训练模型已经学会了边缘、纹理、颜色分布等通用视觉特征我们只需要微调最后几层让它适配果蔬分类。选 ResNet18 而不是 ResNet50主要是计算资源的考量。4200 张图的规模下 ResNet50 的精度优势不明显但训练时间和显存占用却实打实地增加。ResNet18 在单张消费级显卡上可以轻松跑起来而且参数量少也不容易在小数据集上过拟合。这类小规模数据集模型容量不是越高越好够用就行。4.2 训练脚本替换分类头、冻结特征层与完整训练循环import torch import torch.nn as nn import torch.optim as optim from torchvision import models import time # 1. 加载预训练模型并替换分类头 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_features model.fc.in_features model.fc nn.Linear(num_features, 36) # 36 个类别 # 2. 冻结前几层只微调后半段减少过拟合风险 for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False model model.cuda() # 3. 损失函数与优化器 criterion nn.CrossEntropyLoss() optimizer optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.001, momentum0.9, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size8, gamma0.5) # 4. 训练循环 def train_one_epoch(model, loader, criterion, optimizer): model.train() running_loss, correct, total 0.0, 0, 0 for images, labels in loader: images, labels images.cuda(), labels.cuda() optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() return running_loss / total, correct / total # 5. 验证循环 torch.no_grad() def evaluate(model, loader, criterion): model.eval() running_loss, correct, total 0.0, 0, 0 for images, labels in loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) running_loss loss.item() * images.size(0) _, predicted outputs.max(1) total labels.size(0) correct predicted.eq(labels).sum().item() return running_loss / total, correct / total best_acc 0.0 for epoch in range(20): start time.time() train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer) val_loss, val_acc evaluate(model, val_loader, criterion) scheduler.step() print(fEpoch {epoch1:02d} | fTrain Loss: {train_loss:.4f} Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f} Acc: {val_acc:.4f} | fTime: {time.time()-start:.1f}s) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_fruit_veg_model.pth) print(f - 已保存最佳模型 (val_acc{val_acc:.4f}))逻辑说明冻结 layer4 之前的所有层意味着这些层只做前向计算、不参与反向传播训练时更新参数的数量大幅减少过拟合风险随之降低显存占用也更小。fc层前的num_features是从原模型读出来的 512替换成输出 36 的全连接层后模型输出就是 36 类的 logits。参数说明SGD 的momentum0.9是图像分类的经典配置weight_decay1e-4提供 L2 正则。StepLR每 8 个 epoch 把学习率衰减一半让损失在后期稳定收敛。torch.no_grad()是验证时一定要加的否则验证过程会构建计算图、占满显存。训练 20 个 epoch 在这个数据规模下大概几分钟到十几分钟足以看到过拟合趋势。4.3 超参数调节方向什么情况该动哪些参数如果训练后验证准确率卡在某个平台上不去先看学习率。lr0.001对微调是合理的起点但如果你的数据分布和 ImageNet 差异比较大可以尝试lr0.0001更保守地微调。反之如果 loss 下降非常缓慢且准确率没有任何提升可能是学习率过小。batch_size32是显存和梯度稳定性的折中。显存够用可以提到 64梯度更新更平滑但要注意 BN 层的统计量在 batch 太小时会不稳定。如果显存紧张降到 16观察验证集准确率是否出现抖动抖动明显就用 16 配lr0.0005。冻结策略也可以分两阶段先冻结全部特征层只训练 fc 层 5 个 epoch再解冻 layer4 一起微调。这种做法在多类且类别间差异较大的数据集上通常能略微提升精度代价是多花一轮训练时间。5. 避坑记录小样本图像分类最容易翻车的五个问题5.1 标签顺序对不上ImageFolder 的字母序和 labels.json 的文件序现象训练正常跑完验证准确率也不低但画混淆矩阵时发现预测结果明显错位比如模型明明输出「banana」对应的索引人工看图片却是苹果。原因ImageFolder按文件夹名的字母顺序编号apple是第 0 类、banana是第 1 类而labels.json可能是按采集顺序写的第 0 项对应的是 banana。两边都对但顺序没对齐。解决训练前先做一次映射对齐把ImageFolder.class_to_idx打印出来和 json 比对。稳妥的做法是直接在代码里以 json 为准重建一个class_to_idx映射确保 json 里第 N 个类别名和文件夹名一致。5.2 验证集准确率高但实拍新图识别差增强过弱与过拟合的矛盾现象验证集准确率 95%拿手机拍一张放在桌上的香蕉丢进去识别结果却是苹果。原因验证集和训练集同源分布高度一致模型「背下了」训练集的颜色构图特征但没有泛化到真实场景。训练时增强太弱是主因。解决把训练增强里RandomResizedCrop的 scale 参数从默认的(0.08, 1.0)收窄到(0.5, 1.0)避免裁剪过狠导致主体被切掉。同时加入RandomHorizontalFlip和适度的ColorJitter。如果还不理想把测试集换成你自己拍的照片来做最终评估而不是依赖原有的 test 集分数。5.3 loss 在下降但准确率纹丝不动类别不均衡在作祟现象训练 loss 稳步下降但 top-1 准确率一直停在 60% 左右上不去。原因数据集中某几个大类比如香蕉、苹果样本量远大于小类模型倾向于把不确定的样本全部归类到大类里因为这样 loss 总体更低。这就是类别不均衡导致多数类主导梯度更新。解决统计完每类数量后对样本量低于平均线 50% 的类别做针对性处理。简单做法是给损失函数加类别权重。如果训练集样本冗余也可以直接每个小类复制增强后的样本凑到接近均值。5.4 直接 resize 成 224 导致形状信息丢失香蕉变成直线条现象训练集准确率很高但验证集上梨和苹果混淆严重。原因Resize((224, 224))把不同宽高比的原始图强行拉伸成正方形梨的椭圆轮廓和苹果的圆形轮廓在拉长/压扁后特征差异变小。果蔬分类中形状是仅次于颜色的判别特征形变后这个线索直接被破坏了。解决使用Resize(256)CenterCrop(224)的流程保留原始宽高比只裁剪中心区域。如果图片主体不在中心导致裁剪后丢目标可以先Resize(224, 224)并配合训练时的RandomResizedCrop让模型见过各种比例的变体。5.5 显存溢出验证阶段忘了关梯度现象batch_size 设 32 时训练正常但每轮验证跑到一半报 CUDA out of memory。原因验证循环里没有用torch.no_grad()模型前向时照样构建完整计算图梯度缓存把显存撑爆了。训练阶段模型处于train()模式、dropout 和 BN 都开着显存占用本来就高验证再来一次必挂。解决验证函数上必须加torch.no_grad()装饰器或者在循环前手动写with torch.no_grad():。另外model.eval()要放在验证循环最前面它会把 BN 层切到用 running_mean/running_var 的推理模式避免用小 batch 统计量引入噪声。6. 训练完怎么验证效果可视化脚本、混淆矩阵与单类别错误分析6.1 先跑通资源里的 show 脚本用肉眼确认数据集没有脏数据资源里带了show.py这个脚本的功能就是把数据集里的图片直接铺开可视化。训练前跑一遍它不是为了仪式感而是为了确认三件事图片有没有损坏打不开的、标签和画面内容是否一致、有没有同一类里混入明显异类的脏数据。python show.py --root train --num_samples 8 --output grid_preview.png运行后如果发现某个「apple」文件夹里混入了明显是梨的图片或者某些图片是全黑的、模糊到无法辨认的先删掉或者移动到单独的noise/目录再开训。这种脏数据在人工采集的数据集里几乎必然存在不清理的话模型会在这些错标样本上浪费拟合能力。6.2 混淆矩阵找出模型到底把哪两类混在一起画混淆矩阵是每个分类项目跑完后的固定动作。准确率只告诉你模型有多好混淆矩阵告诉你模型哪里不好。果蔬分类里最常见的混淆是洋葱和土豆——两者颜色相近、形状都是扁圆如果模型在这两类之间频繁出错那不是模型不行是这两类本身视觉区分度就低。import numpy as np import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import confusion_matrix all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in test_loader: outputs model(images.cuda()) _, predicted outputs.max(1) all_preds.extend(predicted.cpu().tolist()) all_labels.extend(labels.tolist()) cm confusion_matrix(all_labels, all_preds) class_names train_dataset.classes # ImageFolder 的类别顺序 # 归一化方便看比例 cm_norm cm.astype(float) / (cm.sum(axis1, keepdimsTrue) 1e-8) plt.figure(figsize(18, 16)) sns.heatmap(cm_norm, annotTrue, fmt.2f, cmapBlues, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(Predicted) plt.ylabel(True) plt.xticks(rotation45, haright) plt.yticks(rotation0) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)拿到混淆矩阵后按行看每一行代表真实类别对角线越接近 1 越好非对角线上的值代表被错分到了哪一类。把混淆最大的几对挑出来对着show.py输出的图片仔细看判断是这两类长得太像还是数据标注本身就错了。如果是标注错了把对应图片找出来修正标签比调模型参数更有效。6.3 置信度阈值做拒识生产环境里的最后一道卡控训练出来的模型在测试集上准确率可能到 96%但真实场景里喂进来的图可能根本不属于这 36 类中的任何一类——拍一张键盘模型也会硬分成某个果蔬。这时候要引入置信度阈值softmax 输出的最大概率低于某个阈值就拒绝预测。import torch.nn.functional as F model.eval() threshold 0.75 # 低于这个置信度就拒识 with torch.no_grad(): img load_your_image(real_world_photo.jpg) # 单张推理 img transform_eval(img).unsqueeze(0).cuda() logits model(img) probs F.softmax(logits, dim1) max_prob, pred_idx probs.max(1) max_prob max_prob.item() if max_prob threshold: print(f置信度仅 {max_prob:.2f}拒绝预测请检查输入图片) else: print(f预测类别: {class_names[pred_idx.item()]}置信度 {max_prob:.2f})阈值设多少取决于你业务里「拒识」和「错识」的代价。如果错识后果严重阈值拉高到 0.85 甚至 0.9如果承受得起少量误报但不能漏阈值降到 0.6。拿一小批「非果蔬」图片测试调整比拍脑袋定一个数字靠谱得多。这套流程——跑 show 脚本看数据、训练时盯 train/val 曲线、训练完画混淆矩阵、上线前设拒识阈值——是我每次做小样本分类都要强制走一遍的固定流程。这份果蔬数据集对我来说价值不只是那 4200 张图而是它恰好让你把「拿到数据 → 迁移学习 → 结果诊断 → 工程落地」整条链路完整跑一遍。希望帮到你。本文还有配套的精品资源点击获取