简介面向图像分类任务这份数据集提供已标注的常见果蔬图像覆盖香蕉、苹果、梨、葡萄、橙子、黄瓜、胡萝卜、辣椒、洋葱、土豆等36个类别共约4200张图片。json文件保存了36个类别的名称对应关系图片已预处理可直接作为分类网络的输入。资源包共2000个文件包含1998张jpg图像、1个可视化脚本和1个json配置文件压缩包大小约364.87MB。训练集、测试集、验证集按类别分目录存放同一类图片集中在一个文件夹可被常见深度学习框架的数据加载接口直接读取运行可视化脚本可批量展示样本快速检查图片质量与标注情况。目前已有119人浏览学习特别适合从事果蔬识别、图像分类网络改进或模型对比实验的学生和研究人员。借助划分好的数据子集与可视化脚本使用者可省去大量数据整理时间把精力集中在网络结构与训练调优上对快速验证算法效果很有帮助。1. 4200 张已标注果蔬图够不够一个图像分类模型的起点做图像分类的人第一脚大多踩在公开数据集上。常见那几个动辄几十万张下载要半天训练要整夜跑完才发现自己连数据长什么样都没看清。这个标题里的数据集走的是另一个路子常见果蔬、多类别、约 4200 张、已标注。规模不大但胜在干净、聚焦、能快速闭环。对刚接触图像分类的开发者、做毕设的学生、以及想验证一个分类模型思路的工程师来说它最大的价值不是刷精度而是让你在一天之内走完“数据读取—模型训练—结果评估”的全流程把图像分类那套东西真正跑通。这篇笔记我就按这个数据集的体量把数据怎么组织、模型怎么训、参数怎么设、坑在哪一次说清楚。2. 先说清楚数据集里到底有什么类别构成与标注格式决定你后面的每一步很多人拿到图像分类数据集第一件事就是解压、开训结果训到一半发现类别对不上、图片打不开、标注格式和模型不匹配。这些问题根子都在于没先搞懂数据集的内部结构。图像分类数据集虽然看起来简单就是一堆图片加标签但它的组织方式直接决定你后面能不能用 PyTorch 或 YOLO 系工具顺利跑起来。4200 张的规模意味着什么先说结论这个量级足够训练一个从零开始的浅层 CNN也足够微调一个预训练的 ResNet 或 MobileNet但不足以让你从头训出一个大模型。2.1 类别构成26 类左右是多类别分类的合理区间每类 150 张上下是常见分布“多类别”这个描述在图像分类里通常指 10 类以上、100 类以下。常见果蔬数据集的类别设计一般围绕日常生活比如苹果、香蕉、橙子、番茄、黄瓜、胡萝卜、土豆、洋葱、青椒、生菜这类高频果蔬也可能细化到苹果的不同品种。4200 张如果分到 26 类左右每类大约 150 张出头这是一个很典型的分布——既不会因为单类图片过少导致模型完全学不到特征也不会因为类别太杂导致标注成本失控。这个量级的类别数对模型选型有一个直接影响如果你的分类目标在 10 到 50 类之间那 ResNet18、ResNet50、MobileNetV3 这些经典图像分类模型的输出层改一下就能用完全不需要上 Transformer 结构。如果你拿到手的实际类别数和这个假设差很远比如只有 5 类或者多达 80 类那数据增强策略和模型容量都要跟着调这点在第四章我会展开讲。拿到数据集后第一件事不是写代码而是先做一次抽样摸底。我建议你用下面这个命令把目录结构打出来看看图片文件是不是按类别分文件夹存放的这是图像分类数据集最常见的组织方式# 假设数据集解压在 ./fruits_veg_dataset 目录下 find ./fruits_veg_dataset -maxdepth 2 -type d | sort | head -50 # 统计每个类别文件夹下的图片数量 for dir in ./fruits_veg_dataset/*/; do dir${dir%*/} echo $(basename $dir): $(find $dir -type f | wc -l) done逻辑说明第一个 find 命令只看两层目录确定数据集的根目录下是否直接就是类别文件夹第二个循环对每个类别文件夹计数图片数量。这两个命令跑完你心里就有数了——哪些类别图片多、哪些类别图片少后续在训练时是否需要对少数类做额外处理一目了然。参数说明maxdepth 2 是因为标准结构是“根目录/类别名/图片.jpg”深度为 2。如果你的数据集是“根目录/子目录/类别名/图片.jpg”改成 maxdepth 3。计数时我用的 find 加 wc -l没有用 ls -l 是因为图片文件多的时候 ls 会带出额外信息管道处理也慢。2.2 标注格式文件夹命名分类与 CSV 映射是两种主流方案读取方式完全不同整理得好的图像分类数据集通常有两种标注格式。第一种是文件夹命名即标签也就是 2.1 里说的结构每张图片的位置就隐含了它的类别这种格式在 PyTorch 里用 torchvision.datasets.ImageFolder 一行就能读进来。第二种是图片文件在一个平铺目录里另附一个 CSV 或 JSON 文件记录每张图片的文件名和类别 ID这种格式常见于从数据标注平台导出的结果需要你自己写读取逻辑做映射。数据集标题里提到“已标注”但没有细说标注格式。这个不确定性是你要在项目一开始就确认的先解压看目录结构如果打开就是一堆子文件夹那就是 ImageFolder 风格如果一堆图片加一个 labels.csv那就是映射风格。如果你是从标注平台拿到的原始导出还可能出现每张图对应一个同名 txt 文件的情况那是检测或分割的格式稍作转换也能用在分类上。我一般会建议刚接触图像分类的人优先掌握 CSV 映射格式的处理方式因为实际工程项目里从标注工具导出的数据多数长这样而且它对类别的增删改更灵活import pandas as pd from PIL import Image import os # 假设 labels.csv 有两列filename, label df pd.read_csv(./fruits_veg_dataset/labels.csv) print(df.head()) # 先看一眼列名和内容避免后续读错列 print(df[label].value_counts()) # 统计每个类别的样本数量 # 检查前 5 张图片能否正常打开这一步能提前发现损坏文件 for fname in df[filename].head(5): path os.path.join(./fruits_veg_dataset/images, fname) img Image.open(path) print(fname, img.size, img.mode)逻辑说明这段代码先用 pandas 读入标注文件并输出类别分布确认数据是否均衡然后遍历前几张图片确认路径拼接正确、文件能打开、图像尺寸和通道模式正常。损坏的图片在训练时会导致 DataLoader 报错或产生无效梯度提前检查能省下不少排查时间。参数说明Image.open 之后最好调用 img.verify() 方法做完整校验因为 open 本身是惰性的只读文件头不真正解码像素数据。如果图片文件只有几百字节的残留open 可能不报错verify 才报。另外注意 path 拼接用的是 os.path.join而不是字符串加号因为在 Linux 和 Windows 下分隔符不同join 会自适应。2.3 4200 张数据的数学账训练集能分到多少、模型容量多大合适4200 张的总量按常见的 8:1:1 划分训练集约 3360 张验证集和测试集各约 420 张。3360 张训练图去拟合一个 1000 类的 ImageNet 预训练模型直接全量微调会很容易过拟合。但如果你把预训练模型当作特征提取器只训练最后的全连接分类层这个数据量又显得绰绰有余。这是这个规模的数据集最核心的定位它适合做迁移学习微调不适合从零训练。模型选型上我强烈建议以 ResNet18 或 MobileNetV3 为起点而不是一上来就上 ResNet50 或 EfficientNet-B4。原因很朴素数据量越小模型容量越要克制。ResNet18 参数量约 1100 万ResNet50 约 2500 万在 3000 多张训练图上ResNet50 的过拟合风险明显更高但精度提升却有限。先用小模型跑通流程再逐步加大这是图像分类项目里最稳妥的路径。如果将来自动化标注工具或采集渠道加入数据量扩充到 1 万张以上再换 ResNet50 也不迟。数据集里如果包含同一类果蔬的多个品种或不同成熟度模型需要学习的特征维度会更高那时候 4200 张对某些难分对来说就略显紧张需要靠数据增强和类别权重来弥补。记住一点在数据量不动的前提下调模型结构的边际收益远低于调数据处理和训练策略的收益。3. 用 PyTorch 把 4200 张果蔬图跑进训练管线迁移学习脚本与参数设定数据摸清了接下来就是把数据集真正送进模型训练。这一步我不会给你贴一个几百行的完整工程代码而是拆成“数据加载—模型改造—训练循环—参数配置”四个环节。每个环节单独聚焦你可以直接复制拼起来用。环境默认是 Python 3.9 以上、PyTorch 2.x、GPU 显存 6GB 以上。如果你的机器没有 GPU也能跑只是时间会长很多后面参数部分我会给对应的调整方案。3.1 数据读取与划分用 ImageFolder 搭一个可复用的 DataLoaderPyTorch 的 torchvision.datasets.ImageFolder 是图像分类数据加载的首选工具前提是你的数据集是“根目录/类别名/图片.jpg”结构。如果你的数据集是 CSV 映射格式先做一个预处理脚本把图片按类别移动进对应文件夹或者直接写一个自定义 Dataset 类。对大多数情况来说ImageFolder 足够用而且自带类别名到索引的映射少写不少代码import torch from torch.utils.data import DataLoader, random_split from torchvision import datasets, transforms import os data_dir ./fruits_veg_dataset # 三个关键操作缩放、数据增强、转张量 transform_train transforms.Compose([ transforms.Resize((224, 224)), # 统一尺寸 transforms.RandomHorizontalFlip(p0.5), # 水平翻转增强 transforms.ColorJitter(brightness0.2, contrast0.2), # 亮度/对比度扰动 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) transform_val transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) dataset datasets.ImageFolder(rootdata_dir, transformtransform_train) train_size int(0.8 * len(dataset)) val_size int(0.1 * len(dataset)) test_size len(dataset) - train_size - val_size train_ds, val_ds, test_ds random_split( dataset, [train_size, val_size, test_size] ) # 注意val 和 test 也要套用验证集 transform val_ds.dataset.transform transform_val test_ds.dataset.transform transform_val train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) test_loader DataLoader(test_ds, batch_size32, shuffleFalse, num_workers4) print(f类别数: {len(dataset.classes)}) print(f训练集: {len(train_ds)}, 验证集: {len(val_ds)}, 测试集: {len(test_ds)})逻辑说明transform_train 里的 RandomHorizontalFlip 和 ColorJitter 属于数据增强它们在每次训练迭代时对图片做随机变换相当于帮模型多看一些变体抑制过拟合。Resize 到 224x224 是 ResNet 系列标准输入尺寸Normalize 用的均值方差是 ImageNet 预训练模型的统计值迁移学习时不要改。random_split 是随机切分但切分后要手动把 transform 改成验证集用的 transform否则验证时也会做随机增强导致评估指标不稳定。参数说明batch_size32 是一个比较稳妥的起点显存 6GB 的卡能跑10GB 以上可以考虑 64。num_workers4 在 Windows 上如果报多进程错误改成 2 或 0Windows 的 DataLoader 多进程支持不如 Linux 稳定这是踩出来的经验。随机增强的 p0.5 表示一半概率翻转不要设成 1.0否则所有图片都翻转过模型会误以为“朝左”是某种类别的固定特征。3.2 模型改造把 ResNet18 的分类头换成你的果蔬类别数迁移学习的思路很简单在 ImageNet 上训练好的模型已经学会了通用的纹理、边缘、形状特征你要做的是把最后一层分类输出换成自己的类别数量然后选择性地训练所有层或只训练最后一层。4200 张这个量级我建议两个策略都试一下先只训练分类头看基线准确率再解锁全部层做微调看能提升多少。后者时间成本高但往往能带来 5 到 10 个百分点的提升。import torch.nn as nn import torchvision.models as models def build_model(num_classes, pretrainedTrue, freeze_backboneTrue): # 使用 ResNet18 作为骨干网络 model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT if pretrained else None) in_features model.fc.in_features # ResNet18 这里是 512 # 替换最后一层全连接 model.fc nn.Sequential( nn.Dropout(p0.2), nn.Linear(in_features, num_classes) ) if freeze_backbone: # 冻结所有卷积层参数只训练 fc 层 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True return model model build_model(num_classes26, freeze_backboneTrue) print(model.fc) # 确认分类头结构 # 统计可训练参数量验证冻结是否生效 trainable sum(p.numel() for p in model.parameters() if p.requires_grad) total sum(p.numel() for p in model.parameters()) print(f可训练参数: {trainable / 1e6:.2f}M / 总参数: {total / 1e6:.2f}M)逻辑说明冻结骨干网络后反向传播只更新最后两层训练速度快、需要的数据量少适合第一轮跑基线。Dropout 加在分类头里能进一步抑制过拟合p0.2 是相对保守的设定因为数据量本来就不大。pretrainedTrue 时 PyTorch 新版本要求用 weights 参数而不是 pretrained 参数ResNet18_Weights.DEFAULT 会下载官方在 ImageNet 上的权重。参数说明freeze_backbone 设为 False 就是全量微调。全量微调时学习率要调低1e-4 到 3e-4 之间比只训练分类头用到的 1e-3 小一个量级。原因是预训练权重已经在一个很好的位置学习率太大会把学到的通用特征破坏掉这在迁移学习里叫灾难性遗忘。Dropout 的 p 值不建议在冻结模式下调太高因为分类头本身参数少p0.5 反而可能欠拟合。3.3 训练循环选对损失函数和学习率前面一半功夫不白费多类别图像分类的损失函数没有悬念就是交叉熵。PyTorch 里 CrossEntropyLoss 会自己处理 Softmax你只需要把模型的原始输出和整数标签喂进去。优化器我习惯用 AdamW在图像分类任务上比普通 SGD 收敛更平稳对学习率的敏感度也低一些。学习率调度器用一个简单的 CosineAnnealingLR 就很好它会自动让学习率在训练后期逐步降低帮助收敛到更平稳的极值点。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max20, eta_min1e-6) epochs 20 best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() train_loss running_loss / total train_acc correct / total # 验证集评估 model.eval() val_correct 0 val_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) val_total labels.size(0) val_correct (predicted labels).sum().item() val_acc val_correct / val_total scheduler.step() if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_fruit_model.pth) print(fEpoch {epoch1}/{epochs} | fLoss: {train_loss:.4f} | fTrain Acc: {train_acc:.4f} | fVal Acc: {val_acc:.4f} | fLR: {optimizer.param_groups[0][lr]:.2e})逻辑说明训练循环里有两个细节容易被忽略。第一模型在训练和验证之间要切换 model.train() 和 model.eval()因为 BatchNorm 和 Dropout 在两种模式下的行为不同不切换会导致验证集结果不稳定。第二验证阶段用 torch.no_grad() 关闭梯度计算内存占用小、速度快反向传播在验证时不需要。best_acc 逻辑记录验证集上表现最好的权重后面测试集和实际推理都用它。参数说明T_max20 要和 epochs 保持一致这样学习率从初始值余弦下降到 eta_min。weight_decay1e-4 是 L2 正则化防止权重过大导致过拟合但如果你的分类头只有两层这个值可以调大一点到 1e-3。如果你没有 GPU 或者显存只有 4GB把 batch_size 降到 16epochs 增加到 30学习率降到 5e-4也能收敛只是慢一些。3.4 训练结果的判断基准4200 张数据集跑到什么程度算及格训练结束后你需要一个客观的参照系。纯随机猜是 26 类的 1/26约 3.8%。只训练分类头的迁移学习通常能跑到 85% 以上的验证准确率。全量微调后正常应该能到 92% 到 96%。如果你的验证准确率低于 80%问题大概率不在模型而在数据处理或者标注质量上直接跳到第四章排查。如果到了 96% 以上先别高兴太早要检查是不是数据划分出了问题比如同一个来源的图片被同时分进了训练集和验证集导致评估虚高——这个坑在 4.4 里我详细说。4. 多类别图像分类的避坑清单5 个让新手翻车的典型问题这一章我按“现象 → 原因 → 解决”的顺序写几条最常踩的坑每一条都是自己或身边人真实遇到过的。这些坑对几万张的大数据集影响相对小但在 4200 张这个量级上每一条都可能直接决定你的项目是顺利完成还是卡死两周。4.1 验证集准确率很高但测试集上表现崩塌问题出在随机划分现象验证集准确率到了 94%你高高兴兴去跑测试集结果只有 70%。差别大得离谱。原因随机划分数据时同一个拍摄批次、同一角度、同一光照条件下的图片被同时分到了训练集和验证集验证集参与了训练分布评估结果虚高。这个问题在果蔬数据集上特别明显因为采集时往往是同一批果蔬一起拍摄的。解决按拍摄批次或场景做划分而不是完全随机。如果数据集没有批次信息一个折中办法是先用聚类特征对图片分组再按组划分。import numpy as np from sklearn.cluster import KMeans from PIL import Image # 用简单颜色直方图做分组避免同源图片跨集合 def image_signature(path, bins32): img Image.open(path).convert(RGB).resize((64, 64)) arr np.array(img) # 分通道计算颜色直方图拼成一个特征向量 hist [np.histogram(arr[:, :, c], binsbins, range(0, 256))[0] for c in range(3)] return np.concatenate(hist) # 给每张图片生成签名向量示例只取前 50 张实际要全量 sample_files ... # 你的图片路径列表 signatures np.array([image_signature(f) for f in sample_files[:50]]) kmeans KMeans(n_clusters5, random_state42) groups kmeans.fit_predict(signatures)逻辑说明这个思路是把相似的图片聚到同一组训练集和验证集分别从不同的组里取降低数据泄漏风险。颜色直方图只是最简单的特征表达实际使用时可以用神经网络中间层特征替代效果更好但复杂度也高。对 4200 张图来说颜色直方图基本够用果蔬类别的背景和主色调差异较大聚类出来的分组和拍摄批次往往接近。参数说明bins32 是颜色直方图的粒度越大特征越细但抗噪能力越差n_clusters5 是假设数据约有 5 个拍摄场景你根据实际情况调整。KMeans 的 random_state 固定下来保证结果可复现。4.2 训练损失不降反升伴随大量 NaN学习率太大和坏图片在作怪现象训练到第三个 batchloss 变成 nan然后一路 nan 到底准确率也变成 0。原因一般有两个一是学习率太高导致梯度爆炸二是数据里混入了损坏的图片文件。我在 2.2 里强调过用 verify 检查图片就是为了堵住这个坑。解决先降低学习率试一次如果还是 nan批量检查 DataLoader 里有没有坏图片。from PIL import ImageFile ImageFile.LOAD_TRUNCATED_IMAGES True # 允许加载截断的图片 # 遍历所有图片做完整解码校验 import os for root, dirs, files in os.walk(./fruits_veg_dataset): for fname in files: if fname.lower().endswith((.jpg, .jpeg, .png)): path os.path.join(root, fname) try: img Image.open(path) img.verify() # 验证文件完整性 except Exception as e: print(f损坏文件: {path} - {e})逻辑说明ImageFile.LOAD_TRUNCATED_IMAGES True 是让 PIL 在遇到文件截断时尝试尽力解码避免直接抛异常中断训练。verify 方法会真正读取文件内部结构发现文件头正常但数据损坏的图片。这段遍历脚本会打印所有问题文件路径你核对后决定是删除、替换还是保留。参数说明遍历时用 os.walk 而不是 glob是因为它能递归处理任意深度的目录和 2.1 里的 find 命令互补一个在命令行快速摸底一个在 Python 里彻底检查。4.3 训练集准确率 99%验证集只有 75%经典过拟合增强没到位现象训练集准确率一路飙升到 99%验证集卡在 75% 上下不再动。原因4200 张的数据量不大模型把训练集的背景、颜色分布都背下来了。解决分两步第一步增强数据变换力度第二步给模型结构做瘦身或加强正则化。增强手段我常用的顺序是随机旋转±15 度、随机裁剪ResizedCrop 到 224 再缩放、随机光照扰动。果蔬分类的场景里旋转和裁剪是最贴近真实分布的增强方式因为水果放在桌面上的角度本来就不固定。要注意增强不能过度比如随机旋转到 90 度以上会把苹果和橙子这种圆形水果变成很难辨认的倾斜状态反而增加学习难度。一个常见的配置是 rotation15这种小角度旋转对绝大多数果蔬类别都是安全的。正则化方面除了 3.3 里的 weight_decay 和 Dropout还可以在训练中引入标签平滑label smoothing。原理是把 one-hot 标签中的 1 改成 1-ε剩余 ε 均匀分给所有类别。这样模型不会对训练集的某个样本过于自信等于变相给数据加了噪声对抑制过拟合有实际帮助。PyTorch 里 CrossEntropyLoss 直接支持 label_smoothing 参数设成 0.1 就行。4.4 某些类别准确率异常高、另一些始终为 0类别不平衡和标注噪声现象训练完看分类报告苹果、香蕉这些大类的 F1 分数很高但某个小众类别比如猕猴桃准确率和召回率全是 0。原因可能是这个类别的图片数量太少模型没见过足够的正样本也可能是这个类别的标注本身有问题比如标签错了模型学到的是错误映射。解决先看类别计数如果某类少于 50 张就要考虑用类别权重补偿或做过采样。from sklearn.utils.class_weight import compute_class_weight import numpy as np # 获取所有样本的标签索引 all_labels [dataset.targets[i] for i in range(len(dataset))] # 这行代码假设 dataset.targets 存在ImageFolder 会自动生成 class_weights compute_class_weight( class_weightbalanced, classesnp.unique(all_labels), yall_labels ) weight_tensor torch.tensor(class_weights, dtypetorch.float32).to(device) criterion nn.CrossEntropyLoss(weightweight_tensor)逻辑说明compute_class_weight 会把样本数少的类别权重调高样本数多的类别权重调低这样损失函数对少数类的分类错误更敏感模型在更新参数时会更关注少数类。class_weightbalanced 是自动计算方式它的公式是样本总数除以类别数再除以该类样本数不需要手动指定。参数说明如果你的某个类别样本数极少比如只有 20 张权重补偿的效果也有限这时候最有效的手段是去补充数据或者用这个类别的所有图片做基础变换扩增到至少 100 张。4.5 使用标注工具时踩的坑LabelImg 画完矩形框却导不出分类标签现象你想用标注工具先清洗或扩展这个数据集结果发现某些标注工具导出的是目标检测格式VOC XML、YOLO txt不是分类的文件夹结构。原因图像分类数据集的标注和检测标注是两回事检测需要画框分类只需要给整张图打一个标签。解决分类数据集的标注工作不建议用画框工具直接按文件夹整理图片是最快的配合 labelstudio 或 x-anylabeling 这类支持分类标注的工具逐张打标。提示如果后续你想把这份果蔬分类数据集升级成目标检测数据集再考虑用 LabelImg 或 CVAT 对图片画框。分类和检测是两条路线训练时的损失函数和数据加载方式完全不同不要混着做。5. 模型训完怎么验证才算数混淆矩阵、置信度抽查与少数类复查训练跑完、准确率看着不错别急着收工。图像分类模型的验证不只是看一个总体准确率数字还要搞清楚它在哪里犯错、犯什么错、错误是否集中在某些特定类别上。对 4200 张的果蔬数据集来说最终交付的模型往往用在自动化分拣、商品识别这类场景一个系统性的错误比如把绿番茄全认成青椒比随机几个零星错误严重得多。做两类检查一是汇总层面的混淆矩阵二是个体层面的置信度抽查。import matplotlib.pyplot as plt import itertools from sklearn.metrics import confusion_matrix def plot_confusion_matrix(model, loader, class_names, device): model.eval() all_preds [] all_labels [] with torch.no_grad(): for images, labels in 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) plt.figure(figsize(12, 10)) plt.imshow(cm, interpolationnearest, cmapBlues) plt.title(Confusion Matrix) plt.colorbar() tick_marks range(len(class_names)) plt.xticks(tick_marks, class_names, rotation45, haright) plt.yticks(tick_marks, class_names) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.show() return cm逻辑说明混淆矩阵展示的是“真实类别”和“预测类别”的交叉统计。对角线上的数值越大说明这个类别的识别越可靠对角线外的非零块说明模型在哪些类别之间产生了系统性混淆。果蔬分类里最常见的混淆是外观相近的品类比如番茄和圣女果、青椒和尖椒、橘子和橙子。这类混淆如果占比高可能说明类别定义本身对模型来说区分度不足需要考虑合并类别或者在标签时引入更细的子类层级。参数说明cmapBlues 只是配色方案换成 YlOrRd 或 Greens 都可以不影响数据。rotation45 保证类别名显示不重叠类别名过长时可以改成 90。confusion_matrix 返回的是一个二维 numpy 数组你可以在这个基础上计算每个类别的召回率代码里用 cm[i, i] / cm[i, :].sum() 就能得到第 i 类的召回率那个值比总体准确率更能反映少数类情况。接着做置信度抽查。随机选一批图片让模型给出每个类别的预测概率找到那些预测置信度最高但预测错误的样本这些往往是标注噪声或者图片本身模糊导致的。你把这几个样本拿出来人工看一眼如果图片本身没问题但模型就是信誓旦旦地分错说明这个类别的训练特征没有学到本质需要补充该类的训练图片而不是调模型。这个过程不需要额外写复杂代码在测试循环里把 top-1 概率和真实标签一起保留下来然后用排序找出置信度最高且预测错误的样本即可。至于最后的落地部署4200 张训练出来的模型建议导出成 TorchScript 或 ONNX 格式比直接保存 PyTorch state_dict 更稳定不依赖 Python 环境。验证模型是否真正可用拿几张训练集之外、真实场景里拍的图片跑一次推理看输出概率分布是否合理。这一步能绕过所有可视化指标直接暴露模型在实战中的表现。我自己的习惯是每轮训练完都会保留一份包含验证集上错误样本的 CSV 文件里面记录文件名、真实标签、预测标签和置信度。下次不管是调整数据还是调模型翻开这个文件就能看到上次的问题在哪不用重新跑一遍推理。这算是我做图像分类项目最值得推荐的一个习惯希望帮到你。本文还有配套的精品资源点击获取
