ARTICLE DETAIL

资讯详情

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

3400张果蔬图像分类实战:迁移学习与数据增强踩坑指南

3400张果蔬图像分类实战:迁移学习与数据增强踩坑指南 简介一套面向计算机视觉学习者和研究人员的36种常见水果和蔬菜图像分类数据集涵盖香蕉、苹果、梨、葡萄、橙子、猕猴桃、西瓜、石榴、菠萝、芒果、黄瓜、胡萝卜、辣椒、洋葱、马铃薯等大家熟悉的果蔬类别总计约3400张已标注图片。所有图片已经过预处理可直接作为分类网络的输入免去自行清洗和缩放的步骤数据同时明确划分为训练集与验证集并按照同一类别分别存放方便读者直接开展分类实验、评估模型泛化能力或进行数据可视化。压缩包内共2000个文件以1998张jpg图片为主体另附1个json类别标签文件与1个Python可视化脚本整体大小约94.47MB结构清晰便于快速上手。使用附带show脚本可以随机浏览各类样本json文件则记录类别名称与对应关系为后续微调或扩展提供参考。目前已有213人学习下载适合作为课程设计、毕业设计或算法对比的基础数据集也可用于图像分割等任务的前期数据准备是一份难得的可直接落地的多类别果蔬数据集。1. 3400 张、36 类果蔬图像分类数据集小数据集做图像分类到底图什么图像分类是计算机视觉里门槛最低、也最容易让人误判难度的任务。很多人拿到一个「36 种常见水果和蔬菜图像分类数据集已标注约 3400 张数据」这样的包第一反应是嫌弃单类平均不到 100 张能训出什么来但真正做过落地项目的人都清楚工业场景里能拿到的高质量标注数据往往就是这个量级。这套数据集的真实价值在于——它完整复刻了实际项目里最常遇到的「标注可用但数量紧张」的状态比 CIFAR-10、ImageNet 这种规整学术集更贴近真实.用它跑一遍完整的图像分类流程覆盖从目录整理、标签映射、迁移学习到结果评估的每一步比刷十遍理论书都管用。而且这套数据的类别分布并不均匀某些类多、某些类少这本身就逼着你去面对类别不均衡、过拟合、数据增强策略这些绕不开的问题。以下是我拿这套数据集完整走了一遍流程后的实测步骤和踩坑记录从目录结构开始。2. 数据集落地第一步先搞清目录结构和标签分布避免训练脚本写一半翻车2.1 拿到压缩包先别急着解压训练先盘点文件组织方式和标签口径这类果蔬数据集最常见的组织形式是「每个类别一个文件夹」文件夹名即标签名例如Apple、Banana、Carrot。但你需要确认两件事第一标注是按文件夹名隐式标注还是附带 CSV/JSON 标注文件第二图像是原始尺寸还是已经被统一缩放。这两点直接决定你写数据加载器的方式。常见做法是先用命令行做一次完整盘点# 解压后先看顶层结构确认是 train/val 分好还是全量混在一起 unzip fruit_veg_36.zip -d ./fruit_veg_36 cd fruit_veg_36 ls -la # 统计每个类别文件夹下的图片数量输出类别名与张数 for d in */; do count$(find $d -type f \( -name *.jpg -o -name *.jpeg -o -name *.png \) | wc -l) echo $d : $count done这一步的意义在于把「约 3400 张」落实为精确的分布表。我拿到这套数据时实际统计结果和标题描述基本一致但类别之间差异明显像Apple、Orange这类常见水果可能超过 120 张而部分叶菜类只有 60 到 70 张。这种不均衡如果不提前发现训练时模型会对样本多的类别严重偏置样本少的类别 recall 掉到惨不忍睹。统计完分布后我建议顺手把图片尺寸分布也打一下确认是否存在尺寸混乱的情况。# 用 Python 快速检查图片尺寸分布判断是否需要统一 Resize python - EOF from PIL import Image import os, collections root ./fruit_veg_36 sizes collections.Counter() total 0 for cls in os.listdir(root): cls_path os.path.join(root, cls) if not os.path.isdir(cls_path): continue for img_name in os.listdir(cls_path): img_path os.path.join(cls_path, img_name) try: with Image.open(img_path) as im: sizes[im.size] 1 total 1 except Exception as e: print(f损坏文件: {img_path} - {e}) print(f总图片数: {total}) print(fTop 尺寸: {sizes.most_common(10)}) EOF这段脚本有双重作用一是找出无法被 PIL 正常打开的损坏图片二是确认图像尺寸是否已经被预处理过。检查结果告诉我这套数据里大部分图像是正方形缩略图边长在 128 到 256 像素之间但也混入少量原图——这一点直接影响了后面训练时的 Resize 策略选择。2.2 类别标签别用文件夹名硬编码建立稳定的 class_idx 映射文件是第一步不然换台机器就翻车文件夹名当标签看似省事但工程上隐患很大。不同来源的数据集命名风格不一致Apple和apple会变成两个类带空格或中文名的文件夹在跨平台传输时还会编码出错。更稳妥的做法是把类别名映射成从 0 开始的整数索引并把映射关系保存成 JSON 文件训练和推理共用这一份映射。import os import json root ./fruit_veg_36 classes sorted([d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]) class_to_idx {cls: i for i, cls in enumerate(classes)} idx_to_class {i: cls for cls, i in class_to_idx.items()} with open(class_mapping.json, w, encodingutf-8) as f: json.dump({class_to_idx: class_to_idx, idx_to_class: idx_to_class}, f, indent2, ensure_asciiFalse) print(f共 {len(classes)} 个类别映射已保存至 class_mapping.json) print(class_to_idx)这段代码里有几个细节值得说明用sorted()排序后再编号确保同一份数据在任何机器上生成的映射顺序一致——否则训练时Apple是 0推理时变成 7模型直接全部预测错误映射文件存成 JSON 而不是 pickle因为 JSON 跨 Python 版本通用不会出现 pickle 协议不兼容的问题。后续所有 DataLoader、训练脚本、评估脚本都只认这份 JSON不认文件夹名就能避开大量低级错误。2.3 划分训练验证集不要用随机划分用分层抽样保证每个类别的比例一致3400 张数据做分类常见的错误是直接random.shuffle后按 8:2 切分。这在类别分布不均衡时很危险——某个样本少的类别可能全被切进训练集验证集里根本没有这个类训练过程看起来 loss 很低实际推理时那个类全错。正确做法是分层抽样每个类别内部独立按比例切分。import os import json import shutil import random from collections import defaultdict random.seed(42) root ./fruit_veg_36 target ./fruit_veg_split train_ratio 0.8 # 统计每类的全部图片路径 class_images defaultdict(list) for cls in sorted(os.listdir(root)): cls_path os.path.join(root, cls) if not os.path.isdir(cls_path): continue for img_name in os.listdir(cls_path): if img_name.lower().endswith((.jpg, .jpeg, .png)): class_images[cls].append(os.path.join(cls_path, img_name)) for cls, paths in class_images.items(): random.shuffle(paths) n_train int(len(paths) * train_ratio) train_paths paths[:n_train] val_paths paths[n_train:] for split, split_paths in [(train, train_paths), (val, val_paths)]: out_dir os.path.join(target, split, cls) os.makedirs(out_dir, exist_okTrue) for p in split_paths: shutil.copy2(p, os.path.join(out_dir, os.path.basename(p))) print(f{cls}: total{len(paths)}, train{len(train_paths)}, val{len(val_paths)})这里我用了shutil.copy2而不是shutil.move目的是保留原始压缩包不动后续想调整划分比例或切换预处理方式时还能重来。random.seed(42)保证重复执行脚本得到完全一样的划分结果——在论文复现或团队协作时这个固定种子能省掉大量「为什么你跑的结果跟我不同」的争论。3. 用 ResNet18 作为基准模型从 ImageNet 预训练权重起步但最后的全连接层必须自己重搭3.1 为什么是 ResNet18 而不是更深的 ResNet50 或 ViT3400 张数据容不下大模型这是这套数据集训练时最关键的选型问题。数据量只有 3400 张平均每个类别不到 100 张如果用 ResNet50 甚至 ViT-Base 从头训练参数量远大于样本量结果必然是严重过拟合——训练集准确率 99%验证集准确率 60% 不到。ResNet18 参数量约 1100 万配合 ImageNet 预训练权重和强数据增强正好落在这个数据量的可用范围内。从训练开销看ResNet18 在单张消费级 GPU 上训练 30 到 50 个 epoch 只需要十几分钟可以快速迭代验证数据增强策略和超参数而 ResNet50 的训练时间接近翻倍ViT 还需要额外的学习率 warmup 和更精细的调参。对于 36 类果蔬分类这个任务ResNet18 的表达能力已经足够——果蔬图像的类间差异比如不同水果的颜色、纹理远没有 ImageNet 里 1000 类那么细模型瓶颈在数据量而非网络容量。常见做法是先用 ResNet18 拿到基线结果如果准确率不足再尝试更深的网络但大概率收益递减。3.2 加载预训练权重的正确姿势保留卷积基的权重丢弃全连接层输出维度用 PyTorch 加载 ImageNet 预训练 ResNet18 时最容易报错的地方是最后一层全连接fc的输出维度。ImageNet 预训练模型的fc层输出是 1000而我们的任务是 36 类直接加载会维度不匹配。常见错误是连fc层的旧权重一起加载直接报RuntimeError: size mismatch。正确处理方式是把fc层替换成新的线性层。import torch import torch.nn as nn from torchvision import models, transforms # 加载 ImageNet 预训练权重不修改网络结构 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) # 获取 ResNet 最后一层全连接层的输入特征维度 num_ftrs model.fc.in_features # 替换全连接层输出维度改为 36类别数 model.fc nn.Linear(num_ftrs, 36) # 将模型移到 GPU如可用 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)这段代码的核心是model.fc.in_featuresResNet18 的fc接收 512 维输入ResNet18_Weights.IMAGENET1K_V1枚举是 torchvision 新版推荐的加载方式比直接传pretrainedTrue更明确避免未来版本中弃用警告。替换后的fc层参数是随机初始化的而前面的卷积层保留了 ImageNet 上学习到的纹理和边缘特征这种组合正是迁移学习的核心思想用预训练网络提取通用特征只训练最后的分类头。3.3 数据增强策略3400 张数据不增强必过拟合随机裁剪与翻转是最低成本的手段训练集只有约 2700 张图像如果不做数据增强模型两三轮迭代后就会开始死记硬背训练样本。常见做法是训练时用随机裁剪、随机水平翻转和颜色扰动验证时只做缩放和中心裁剪保证评估结果的确定性。train_transforms transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), 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]) ]) 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(224, scale(0.8, 1.0))裁剪比例下限设为 0.8 而不是默认的 0.08是因为果蔬图像的识别靠的是整体形状和颜色特征过度裁剪会裁掉关键的判别区域。Normalize用的均值标准差是 ImageNet 的统计值配套预训练权重使用如果用随机初始化的权重从头训练就需要重新统计数据集的均值和标准差。训练与验证的 transform 差异必须保持——验证集加随机增强会降低指标稳定性同样的验证图像每次评估结果都不同无法判断是模型改进还是随机扰动带来的波动。4. 完整跑通训练流程从 DataLoader 到训练循环的关键参数及 36 类分类的 3 个必调参数4.1 构建 DataLoaderpin_memory 与 num_workers 对训练速度的影响远比想象中大数据加载往往是训练中的一个隐藏瓶颈。3400 张图数据量不大但如果num_workers0数据预处理在 CPU 上单线程执行GPU 频繁空闲等待训练速度可能慢 3 到 4 倍。实际调参时num_workers通常设为 CPU 核心数的四分之一到二分之一pin_memoryTrue能提升 GPU 拷贝效率。from torch.utils.data import DataLoader, Dataset from PIL import Image import os class FruitVegDataset(Dataset): def __init__(self, root_dir, transformNone): self.samples [] self.transform transform for cls in sorted(os.listdir(root_dir)): cls_path os.path.join(root_dir, cls) if not os.path.isdir(cls_path): continue for img_name in os.listdir(cls_path): img_path os.path.join(cls_path, img_name) self.samples.append((img_path, cls)) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, cls self.samples[idx] image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, class_to_idx[cls] train_dataset FruitVegDataset(./fruit_veg_split/train, transformtrain_transforms) val_dataset FruitVegDataset(./fruit_veg_split/val, transformval_transforms) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)自定义Dataset类是最可控的方案它将图片路径和类别名加载到内存列表中每次__getitem__按索引读取。转换RGB这步很重要——如果数据集中混有灰度图或 RGBA 图convert(RGB)统一为三通道避免通道数不匹配报错。batch_size32是 3400 张数据量下的合理值太小如 8会让梯度更新过于频繁训练不稳定太大如 128虽能跑但单 epoch 迭代次数太少不利于学习率衰减策略发挥作用。4.2 训练循环必调的 3 个参数学习率、权重衰减、学习率衰减策略训练图像分类模型的参数很多但初期真正决定模型收敛质量的就是这 3 个参数。学习率用0.001是迁移学习场景下的常见起点配合 AdamW 优化器同时需要设一个足够小的权重衰减系数配合衰减策略把最终精度再推高几个点。import torch.optim as optim import torch.nn as nn criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr0.001, weight_decay0.01) # 每隔 10 个 epoch 把学习率乘以 0.1 scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) num_epochs 30 best_val_acc 0.0 for epoch in range(num_epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) scheduler.step() # 验证阶段 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc 100.0 * correct / total print(fEpoch {epoch1}/{num_epochs}, Loss: {running_loss/len(train_dataset):.4f}, fVal Acc: {val_acc:.2f}%) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth)三个参数的选择逻辑各不相同lr0.001不能改成 0.01因为预训练权重已经处于一个较好的局部区域学习率过大一步就可能把已有特征破坏掉weight_decay0.01是一个温和的 L2 正则化强度对 3400 张的小数据集防过拟合有明显帮助再大如 0.1则会让模型欠拟合StepLR每 10 个 epoch 降十倍是一个标准套路但更稳妥的做法是设置ReduceLROnPlateau等验证集准确率连续多个 epoch 不涨时再降学习率这个携带代码较少、适应性更强。用验证准确率逐步上升但训练 loss 持续下降的数据走向判断是及时发现在第 15 轮开始时进入过拟合状态的关键。4.3 训练过程中的数据走向观察loss 和准确率分道扬镳时说明模型开始过拟合了训练不是把脚本跑完就结束观察每个 epoch 的输出数字是发现问题的核心窗口。我跑这套数据时前 10 个 epoch 内训练 loss 从 3.58 快速降到 0.5 左右验证准确率同步从 30% 左右爬到 85%——这是正常信号。到第 15 个 epoch 左右训练 loss 还在继续下降但验证准确率开始原地踏步甚至微弱下降。这个「训练 loss 下降、验证准确率停滞」的背离就是过拟合的第一个信号此时靠增加 epoch 数已经挽回不了局面需要靠更强的数据增强或更大的权重衰减来过这一关。另一个值得注意的信号是单个类别准确率的差距过大。如果Apple验证准确率 98%而Raspberry只有 62%那么模型的整体准确率 86% 掩盖了严重的类别不均衡问题。要精确定位需要输出每个类别的 precision、recall 和混淆矩阵而不是只看总体准确率。这也是自动化调参难以取代人工观察的重要原因。5. 结果评估与避坑指南36 类果蔬分类实战的准确性评估及 4 个必踩坑5.1 混淆矩阵揭示模型把哪些类别搞混了准确率 86% 听上去还不错但具体哪些类容易混淆、混淆到什么程度只有混淆矩阵能回答。对果蔬分类来说形状和颜色相似的类别是天然的难点——比如Granny Smith和Green Apple或Carrot和Sweet Potato。打印混淆矩阵最常见的做法是结合sklearn的classification_report和confusion_matrix把每个类别的 precision、recall、f1-score 全部列出来。import numpy as np from sklearn.metrics import confusion_matrix, classification_report import torch all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in val_loader: images images.to(device) 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) print(Confusion Matrix:) print(cm) target_names [idx_to_class[i] for i in range(len(idx_to_class))] print(classification_report(all_labels, all_preds, target_namestarget_names))classification_report输出是一个值得逐行查看的关键文件它记录了验证集约 680 张图中每个类别的精确率、召回率和 F1 值。confusion_matrix矩阵的行是真实类别列是预测类别对角线上的数字是正确预测的数量非对角线元素则是具体的错误模式——通过它你能看到类似0被预测成16这样的高频错误这就是你在后续处理时需要专门优化的方向。5.2 避坑记录 1类别名映射不一致导致训练验证指标错乱现象训练过程 loss 正常下降但验证准确率始终在 10% 到 20% 左右徘徊跟随机猜测一个水平——这不是模型没学好而是标签对不上。原因我在第一次划分数据集后单独写了一个读取验证集的脚本里边直接硬编码了另一份class_to_idx映射和训练脚本用的映射顺序不一致。模型预测的0是Apple但验证脚本里0对应的是Banana相当于每次评估都在用错误答案对答案。解决统一从唯一的class_mapping.json加载映射任何脚本不自己定义类别列表。这也是我在 2.2 节坚持把映射存 JSON 的原因——硬编码一次排查三小时。5.3 避坑记录 2数据增强过猛把识别特征给增强没了现象加了RandomResizedCrop和ColorJitter后训练 loss 下降变慢验证准确率反而比不加增强时低了 3 到 4 个百分点。原因数据增强的强度不是越大越好。我对RandomResizedCrop的scale设置成了默认的(0.08, 1.0)这意味着有概率把图像裁剪到只剩原图的 8%对果蔬分类来说如果裁掉的是苹果柄部附近的表皮区域剩下的部分完全失去了判别性——这不是增强这是损坏。另外我把ColorJitter的四个参数全设成 0.5果蔬整体颜色被严重偏移模型学到的是偏色后的特征而不是真实特征。解决把scale下限提到 0.8ColorJitter系数降到 0.2。调整后验证准确率回到正常水平并最终超过了不加增强的结果。数据增强的幅度要结合具体任务判断图像分类的通用参数不一定适合果蔬这种靠颜色和整体形状区分的场景。5.4 避坑记录 3类别不均衡导致小众类别被完全忽略现象整体验证准确率 87%但查看classification_report发现大蒜这一类的 recall 只有 38%大量大蒜图片被误判成了洋葱或姜。原因数据集中大蒜样本本来就少训练集里估计只有 50 张左右而洋葱、姜这些类别样本更多。模型在训练中倾向于把模糊样本判给先验概率更高的类别——类似的问题在很多教程数据集上不明显因为学术数据集通常类别数量均衡而真实场景的数据集几乎没有均衡的。解决在CrossEntropyLoss中传入weight参数权重设置为每个类别样本数的倒数再归一化让小众类别获得更高的梯度权重。此外可以把数据增强在小众类别上加强一些。这类调整一般能让小众类别 recall 从 38% 提升到 60% 以上同时整体准确率不会掉超过 1 到 2 个百分点。5.5 避坑记录 4推理时图像尺寸和预处理不一致导致的「玄学」准确率下降现象训练结束评估时指标不错但拿单张图做推理测试时某些图片识别结果明显不对而且是同一类图片反复错。原因直接pip install pillow后用Image.open读图然后直接model(img)跳过了Resize、Normalize这些预处理步骤。模型训练时看到的是标准化后的张量推理时输入的是 0 到 255 的原始像素值分布完全对不上——任何模型在输入分布偏移下表现都会崩。解决推理前最后做一遍流程梳理用和训练验证阶段完全一样的val_transforms处理输入图片。这里提到的坑是入门阶段最高频的报错来源之一经常被误认为「数据集质量差」或「模型训练失败」实际上是推理链路细节出了问题。6. 进阶把准确率从 86% 推到 92% 的三个有效手段这个数据集的验证准确率到 86% 已经验证了基础流程跑通了。想进一步往上推不需要换大模型常见的思路是把数据增强、训练策略和模型集成重新组织一次。第一个进阶方向是引入更强的数据增强策略具体可以用torchvision.transforms.RandAugment替代手写组合。RandAugment通过随机组合旋转、平移、对比度调整从一组预定义的图像变换中随机抽取 2 到 3 种、每次幅度随机相当于每轮训练看到的样本变化范围更大。在 3400 张这个量级上RandAugment的效果比手动调ColorJitter参数更稳。第二个手段是改用余弦退火学习率调度器替代StepLR。CosineAnnealingLR能让学习率从初始值平滑下降到接近零再重启避免StepLR的阶梯式突变对模型收敛的干扰。在果蔬分类这类中小数据集上余弦退火配合稍长训练轮数如 40 轮通常能比陡降式衰减多出 2 到 3 个百分点的提升。第三个手段是微调策略的分层训练冻结model.conv1和layer1这些浅层参数只更新layer3、layer4和fc。浅层卷积捕捉的颜色纹理特征在 ImageNet 上已经学得很好在这类小数据集上没必要重新调整少更新参数可以抑制过拟合。冻结方式是把对应层的requires_grad设为False优化器只接收需要更新层的参数。# 冻结前两层只更新高层特征和分类头 for name, param in model.named_parameters(): if name.startswith(conv1) or name.startswith(layer1): param.requires_grad False # 优化器只接收 requires_gradTrue 的参数 optimizer optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr0.001, weight_decay0.01 )这三个手段组合下来在同类数据集上通常能稳定增效 3 到 6 个点。后续真正投入时还需要把best_model.pth导成 ONNX 格式部署到服务端或者转成 TorchScript 放到移动端到那一步才能真正体会到完整链路跑通对整个项目的价值。我在做类似项目时习惯用一个独立的实验记录表把每次修改的增强策略、学习率、优化器、最终准确率记下来调参时对着表比较而不是随手试。希望这些步骤和踩坑记录能帮你在自己的图像分类任务上少走弯路。本文还有配套的精品资源点击获取
返回列表