首页 / 资讯中心 / 文章详情

蘑菇图像分类实战:PyTorch ImageFolder加载与迁移学习避坑指南

蘑菇图像分类实战:PyTorch ImageFolder加载与迁移学习避坑指南 ★ FEATURED ARTICLE
简介面向蘑菇图像识别与计算机视觉分类任务这份数据集将蘑菇图片按训练集与测试集划分好并附带类别字典文件省去自行整理标注的麻烦适合深度学习入门者以及使用YOLOv5、CNN等框架开展图像分类实验的开发者。资源包共2000个文件以1998张JPG图片为主体另含1个Python可视化脚本和1个JSON类别字典压缩包约97.67MB目录结构清晰训练集与测试集分别存放便于直接加载。目前已有1048人浏览/学习过该资源。训练集含9600张图片、测试集含2400张图片可满足图像分类网络的基本数据需求JSON文件定义了类别信息show脚本可帮助快速预览样本涵盖姬松茸、牛肝菌、Cortinarius等多个类别适合直接用于模型训练、迁移学习或YOLOv5分类任务减少数据预处理工作量。1. 你的第一份蘑菇数据集其实是来教你处理「真实分类麻烦」的做图像分类的开发者第一次拿到这份 12 种蘑菇图像识别数据集时多半会愣一下打开文件夹看到的不是一堆散落的 jpg而是一个层级清晰的目录——train、val、test 各占一块每个子目录按类别命名。旁边还躺着一个类别字典文件把标签转成索引的映射关系写得明明白白。这种组织方式恰恰是工业界图像分类项目最标准的输入形态也是很多教程里很少教你的部分数据划分已替你完成模型训练的第一步是「正确读取」而不是「重新整理」。如果你正准备入门分类任务、做迁移学习或需要一份随时能跑通 baseline 的数据集这份蘑菇数据会省掉大量清洗和打标的体力活。真正值得关注的不是「蘑菇有多少种」而是「这份数据怎么被组织、加载时容易在哪儿出错、什么参数能把它训练到可用」。2. 数据集内部结构文件夹划分和类别字典文件到底怎么组织2.1 文件夹保存的目录约定train/val/test 与类名子目录我拿到手的第一件事是用 tree 命令把目录全貌打出来确认划分逻辑。这个数据集常见的目录长这样mushroom_dataset/ ├── train/ │ ├── amanita_muscaria/ │ ├── boletus_edulis/ │ ├── cantharellus_cibarius/ │ └── ... ├── val/ │ ├── amanita_muscaria/ │ ├── boletus_edulis/ │ └── ... ├── test/ │ ├── amanita_muscaria/ │ └── ... └── categories.json每个类别一个子目录文件名可能是原始编号或带前缀的哈希名这不重要关键在层级——PyTorch 的torchvision.datasets.ImageFolder会直接把这个目录结构读成带标签的数据集目录名就是类别名。val 和 test 同时存在意味着你不需要自己再切一份验证集直接按给定的划分训练、调参、最终评估即可。和「一个 CSV 写所有图片路径 标签」的平铺形式相比文件夹保存对新手最大的好处是「可见」打开目录就能数出每类多少张能直接肉眼检查某张图是否混入了错误类别。而你带着 val/test 一起用得出的精度才值得写进报告里。实际项目里我见过不少人只跑 train 就宣称精度多高结果拿到 val 上一落千丈——划分文件的意义就是逼你用统一口径评估。2.2 类别字典文件从 txt 到 json 的两种常见格式类别字典文件通常有两种长相一种是朴素的 txt一行一个类名另一种是 json直接给出类名到整数索引的映射。做法上json 更友好{ amanita_muscaria: 0, boletus_edulis: 1, cantharellus_cibarius: 2 }txt 的加载则要自己排序保证顺序和文件夹名一致with open(categories.txt, r) as f: classes [line.strip() for line in f.readlines() if line.strip()] class_to_idx {name: i for i, name in enumerate(classes)}这里最大的风险是「顺序错位」——txt 加载后如果排序逻辑和 ImageFolder 的sortTrue默认行为不一致类别索引就会整体漂移训练出来的模型预测结果全对不上号。我习惯在读数据后立刻打印一次dataset.class_to_idx跟字典文件里的映射逐行比对肉眼确认再开工。2.3 数据量、类平衡与标签质量动手前先做这三项体检训练之前先写一个几十行的小脚本把数据的基本面摸清楚。至少要检查三点每类图片数量、图尺寸分布、有没有打不开的坏图。import os from collections import Counter from PIL import Image root mushroom_dataset/train counts Counter() sizes Counter() for cls in os.listdir(root): cls_dir os.path.join(root, cls) if not os.path.isdir(cls_dir): continue for img_name in os.listdir(cls_dir): counts[cls] 1 try: with Image.open(os.path.join(cls_dir, img_name)) as im: sizes[f{im.width}x{im.height}] 1 except Exception: print(fcorrupt: {os.path.join(cls_dir, img_name)}) print(counts) print(sizes.most_common(10))这一步会暴露两个问题一是类不平衡比如某种蘑菇只有 100 张另一种有 800 张二是图尺寸不统一有些是 224x224有些是 800x600。这两点直接决定了后面数据增强和 Resize 的参数选择。类不平衡严重的得在采样器上做文章尺寸差距大的ImageFolder 默认的 Resize 会硬拉伸导致蘑菇形态失真。先体检再训练这是我做数据集项目时雷打不动的第一步。3. 用 PyTorch 把蘑菇数据集跑起来从目录到 DataLoader 的最小实现3.1 ImageFolder 直接读文件夹三行代码加载训练集PyTorch 的 ImageFolder 为这种文件夹结构量身定做不需要自己写 Dataset 类。最小可用代码from torchvision import datasets, transforms train_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_dataset datasets.ImageFolder(mushroom_dataset/train, transformtrain_transforms) val_dataset datasets.ImageFolder(mushroom_dataset/val, transformtrain_transforms)这段代码里Resize((224, 224))把图固定到正方形以适配标准分类网络输入ToTensor()把 HWC 转 CHW 并归一到 [0,1]Normalize用的是 ImageNet 统计值这是迁移学习时最常见的选择。如果你打算从头训练网络可以把均值和方差换成自己数据集的统计值或者干脆不归一化但迁移学习场景下统一 ImageNet 参数是惯例效果最稳。同样一个变换对象同时用在 train 和 val 上是刻意为之验证集不参与数据增强只做尺寸调整和归一化保证评估口径一致。这点太多人在第一次跑分类时搞混直接在验证集上开随机翻转导致精度虚高且无法复现。3.2 类别字典与标签映射确保训练脚本不张冠李戴ImageFolder 会自动扫描子目录名生成class_to_idx这一步通常是可信的但前提是目录名不要出现中文、空格或特殊字符。之前在某容器环境里我见过一次加载后分类数变 14 的怪事排查发现是有两个隐藏目录._amanita_muscaria被 macOS 归档时带了出来ImageFolder 把它们也当成了类别。所以目录干净是第一步加载后立即验证是第二步print(train_dataset.classes) print(train_dataset.class_to_idx) assert len(train_dataset.classes) 12, 类别数不是 12检查目录下是否混入非图片文件夹这是便宜且有效的防错手段。类别字典文件此时不是摆设——它应该与class_to_idx完全一致。如果数据集给的是 categories.txt而 ImageFolder 生成的顺序是按目录名字母序两者可能不一致直接用数据集给的字典文件覆盖import json with open(mushroom_dataset/categories.json, r) as f: custom_map json.load(f) # 用自定义映射重排保证与数据原作者划分一致 train_dataset.class_to_idx custom_map train_dataset.classes list(custom_map.keys())此操作的代价是后续所有预测结果都要以 custom_map 为准而不是 ImageFolder 的默认顺序。我通常会再强制把dataset.targets按新映射重新编码一遍避免 sampler 与映射不一致时标签错位。3.3 数据增强参数怎么设旋转、翻转、归一化的推荐值蘑菇图像和猫狗照片不同蘑菇往往从俯视或侧视角度拍摄分类特征集中在菌盖形状、颜色、菌褶纹理上。增强策略要贴合这个生物特征train_transforms transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.05), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])三点说明一是RandomResizedCrop替换固定 Resize让模型看到不同尺度和裁剪位置的蘑菇增强对拍摄远近变化的鲁棒性二是RandomRotation只转 15 度不是 90 度——蘑菇没有强方向性但转太多会让菌盖形状失真三是色相抖动设得很小0.05因为蘑菇颜色是重要分类线索色相大幅扰动等于故意制造错误样本。这些参数不是玄学在后续验证里能看到实际影响增强过猛val 精度波动增大增强过弱train 精度高但外推差。4. 训练与验证一版能复现的分类脚本和三个必调参数4.1 分 batch 训练与验证循环一个可抄的模板直接上一个能完整跑通的训练脚本骨架注释里写明每个环节的作用import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import models train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 12) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) best_acc 0.0 for epoch in range(30): model.train() running_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) model.eval() correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) val_acc correct / total print(fepoch {epoch1}: loss{running_loss/len(train_dataset):.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_mushroom_model.pth)整个循环分了「训练态」和「评估态」两段用model.train()和model.eval()显式切换。训练态开启 Dropout 和 BatchNorm 统计更新评估态冻结这些随机因素保证 val 指标真实。CrossEntropyLoss内部自带 Softmax不需要在模型最后一层重复加。保存最佳权重用 val_acc 作门槛而不是看最后一个 epoch——后面 20 轮过拟合的风险很高只存最优能留一条后悔药。4.2 学习率、batch size 和图像尺寸初始值怎么定三个参数里学习率是最容易翻车的迁移学习用1e-4起步是稳妥的从头训练则可以放到1e-3。这里给一组可复现的初始值建议参数初始值调整依据学习率 lr1e-4迁移学习/ 1e-3从零训练loss 震荡太剧烈则减半batch size32显存不够先降到 16batch 太小建议同步调低 lr输入尺寸224x224数据集中有大量小图时改 192提速明显优化器AdamW比 SGD 收敛稳追求极致精度可换 SGDmomentum训练轮数30看 val_acc 是否连续 5 轮不涨提前停batch size 决定梯度估计的噪声水平。小 batch 噪声大反而有正则效果但 Adam 在小 batch 下的稳定性不如大 batch。我一般先固定 224 和 32跑 5 个 epoch 看 loss 曲线如果 loss 在第 3 轮还在 2.4 以上不动就把 lr 调大一倍如果第 2 轮就开始震荡就减半。不要凭感觉改多个参数一次只动一个。4.3 结果保存最佳权重、日志与类别名对齐训练收尾时要把三样东西一起存下来模型权重、训练日志、类别映射。权重文件只存state_dict不存完整模型便于换架构时复用日志可以用 pandas 记成 CSV类别映射直接沿用 categories.json。恢复推理时的核心代码如下import json checkpoint torch.load(best_mushroom_model.pth, map_locationcpu) model.load_state_dict(checkpoint) with open(mushroom_dataset/categories.json, r) as f: idx_to_class {v: k for k, v in json.load(f).items()} # 推理时把输出索引映射回蘑菇名称 logits model(sample_tensor) pred_idx logits.argmax(dim1).item() pred_name idx_to_class[pred_idx]这一步看着琐碎但实际落地时八成以上的错误出现在这里模型在 GPU 上训练保存推理时忘记map_locationcpu导致加载失败或者类别映射用的是 ImageFolder 自动顺序而不是数据集的 json导致「预测成第 5 类」其实是第 7 类。权重、日志、映射三者捆绑保存是我在每个分类项目里固化的习惯能少走一大截弯路。5. 蘑菇数据集实战避坑加载报错、类不平衡与易混淆对的排查5.1 ImageFolder 分类数对不上先查文件夹层级现象train_dataset.classes打印出来有 13 个不是预期的 12或者训练时 loss 正常但 val_acc 在第 10 轮后停滞在异常低位。原因最常见是目录里混入了非图片文件或隐藏文件夹。macOS 的._前缀文件、编辑器生成的.DS_Store、解压不完全产生的临时目录都会被 ImageFolder 当成独立类别。另外一种情况是某个类别子目录下还有一层子目录导致该类别被嵌套扫描出多个子类。解决先跑一个打印脚本把classes和每个类别的样本数输出逐个比对。清除隐藏文件和多余层级后重新加载find mushroom_dataset/train -type f -name ._* -delete find mushroom_dataset -type d -name __MACOSX -exec rm -rf {} 删除后重新打印 classes确认正好 12 类再开工。这个检查成本极低收益极高。5.2 图片损坏导致训练中断用 PIL 做一次批量体检现象训练跑到第 7 轮某个 batch 突然抛OSError: image file is truncated程序直接崩掉前 6 轮白跑。原因数据集里存在损坏或不完整的 JPEG——可能是采集时文件没写完也可能是压缩传输时丢字节。这些坏文件在最开始的体检脚本里没有被捕获因为Image.open()默认惰性加载只有真正读取像素数据时才会抛错。解决用ImageFile.LOAD_TRUNCATED_IMAGES True加显式加载做全量体检。更彻底的方案是直接把损坏文件移出数据集from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES True bad_images [] for root, dirs, files in os.walk(mushroom_dataset): for fname in files: fpath os.path.join(root, fname) try: with Image.open(fpath) as im: im.load() # 强制读像素 except Exception: bad_images.append(fpath) for fp in bad_images: os.replace(fp, fp .bad) print(ffound {len(bad_images)} corrupt files)处理完坏图后重新统计每类数量避免误删造成新的类不平衡。5.3 类不平衡12 个类数量差异大时怎么办现象训练时整体 val_acc 到了 88%但打印每个类别的 precision 发现「某种少见蘑菇」只有 40% 多大头全被高频类拉高了。原因数据集中某些类样本少模型在 CrossEntropy 的驱动下倾向于把所有不确定样本判给高频类。文件夹结构本身不会解决这个问题需要从采样和损失函数入手。解决第一档用WeightedRandomSampler按样本量反比加权采样让小类在每个 epoch 里被抽到足够多次from torch.utils.data import WeightedRandomSampler targets train_dataset.targets class_counts torch.bincount(torch.tensor(targets)) weights 1.0 / class_counts.float() sample_weights weights[targets] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler)注意使用 sampler 时要把 DataLoader 的shuffle设为 False两者冲突。第二档如果重采样后小类仍不稳再考虑Focal Loss它会让模型把注意力集中在难分样本上。类不平衡不一定会让整体精度崩盘但它一定会在测试阶段反咬一口——12 类里总有几种蘑菇外观接近小类识别率差意味着一张照片会被误报成「另一类毒蘑菇」这个风险在真实场景里不可接受。5.4 毒蘑菇与可食蘑菇易混淆看 top-2 混淆对数现象val_acc 看着不错但打开混淆矩阵发现大部分错误集中在「某两种白色蘑菇」互相误判。原因蘑菇分类的难点之一就是同类近似——同属不同种之间菌盖颜色、菌柄形态高度相似。CrossEntropy 只会给一个正确答案不会告诉你「第二可能的答案也很接近」而这些近邻信息才是实际部署时最有价值的信号。解决推理阶段不仅取argmax同时取 Top-2 概率输出probs torch.softmax(logits, dim1) top2_probs, top2_indices torch.topk(probs, 2, dim1) for i in range(len(top2_indices)): top1 idx_to_class[top2_indices[i, 0].item()] top2 idx_to_class[top2_indices[i, 1].item()] if top2_probs[i, 1].item() 0.3: print(fambiguous: {top1} vs {top2})通过这个脚本统计全验证集能快速定位哪些类对是模型的高频混淆源。混淆严重的类对就要考虑在训练时单独为它们多采集样本或者在后处理时引入「拒识」逻辑——当 Top-1 和 Top-2 概率都接近时输出「无法判断」而不是硬给一个答案。5.5 验证精度远低于训练精度检查归一化与是否误用训练增强现象train_acc 到了 97%val_acc 只有 72%差距悬殊。原因最大的嫌疑是验证集误用了训练增强。比如我把RandomResizedCrop也放进了 val 的 transforms 里导致验证时每次评估结果都不一样数值偏低且不可复现。其次是归一化参数不匹配——如果训练时用了 ImageNet 的 mean/std推理或验证时用了自己算的统计值激活分布直接错位。解决验证集只用Resize((224,224)) ToTensor() Normalize三段不要任何随机增强。如果确认变换没问题再看是不是模型已经过拟合训练集小、网络容量大学到最后就是把训练图硬背下来。此时先加 weight decay 到1e-3再把轮数减到 20配合早停观察 val_acc 的拐点。过拟合和增强误用是「训练高验证低」的两大元凶排查顺序按「验证集变换 → 归一化 → 模型容量」来走通常前两步就能找到问题。6. 让蘑菇分类更准的进阶技巧迁移学习、置信度阈值与混淆矩阵在 baseline 跑通之后想进一步提升识别能力三个方向性价比最高换更大的预训练骨干网络、上置信度阈值、用混淆矩阵做定量分析。迁移学习可以直接把 resnet18 换成 resnet50 或 efficientnet-b0改动只有三行model models.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 12)换骨干后一般要保持同样的学习率但把训练轮数降到 20 以内。resnet50 的参数量接近 resnet18 的两倍在小数据集上更容易过拟合所以早停要盯得紧一点。置信度阈值是落地时最实用的一个手段。实际照片的拍摄条件千差万别总会出现「模型被迫从四个模糊类里硬挑一个」的时候。我给推理脚本加上置信度门槛低于 0.6 的直接输出「不确定」prob, idx torch.max(torch.softmax(logits, dim1), dim1) if prob.item() 0.6: print(reject: confidence too low, recommend human check) else: print(fpredict: {idx_to_class[idx.item()]}, confidence: {prob.item():.2f})这个阈值的确定方法很简单在验证集上扫一遍 0.5 到 0.9看精度和拒绝率怎么变化画出一条 trade-off 曲线选取业务上能接受的拒绝率对应的阈值。最后一步是混淆矩阵。用 sklearn 在验证集上跑一遍把 FP 和 FN 最大的几个类对列出来我会逐个翻几张被分错的原始图确认是「外观确实相近」还是「数据集标错了」。很多时候标注错误比模型错误更常见——这一点只有看了具体图片才能发现。从这份 12 种蘑菇数据集里我学到的第一课是好的数据集不只是图片多而是划分合理、字典清晰、坑可以被预期。在你动手换更大的模型之前先把数据体检做扎实比任何结构升级都有效。希望这份笔记能让你少走一点我走过的弯路也欢迎你把跑出来的混淆矩阵翻出来看看——那是数据质量最诚实的一面镜子。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助?
咨询建站