简介这是一份面向深度学习图像识别任务的十一种水果分类数据集覆盖苹果、鳄梨、蓝莓、辣椒、樱桃、猕猴桃、芒果、橙子、岩瓜、草莓、小麦等常见类别适合入门图像分类、训练卷积神经网络以及验证不同网络结构的场景。数据严格按类别分文件夹整理训练集约两千五百余张、测试集约六百余张无需额外标注解压即可直接用于当前主流深度学习框架。压缩包约八百五十余兆共两千个文件图像主体以JPEG格式为主另含少量PNG、WebP、BMP图片并附带类别字典和可视化脚本便于查看类别标签分布和抽样检查图像质量。目前已有超过一千六百人次浏览学习下载后可直接以文件夹名称作为标签省去繁琐的数据清洗流程。对于需要高质量水果图像数据、想专注于模型设计而非数据整理的开发者这份资源能有效提升前期准备效率。1. 11类水果分类数据集为什么分类任务比检测更适合入门做深度学习图像识别很多人一上来就奔着 YOLO 系列去想跑目标检测。但如果你是第一次接触视觉项目我更建议从分类任务起步。这个 11 种水果分类数据集就是一个典型的 11 分类场景苹果、香蕉、葡萄、橙子这类常见水果每类图片数量够用背景不算太干净但也没到故意刁难的程度拿来练手 CNN 或者迁移学习正合适。这个数据集的定位很明确不追求检测框只做整图分类。你给它一张水果图片它输出这个水果属于哪一类。相比检测任务分类不需要标注框只需要文件夹级别的标签数据准备成本低训练收敛也快。适合刚接触深度学习的读者跑通全流程也适合想验证 ResNet、EfficientNet 这类模型在小规模数据上表现的老手做基准测试。我自己第一次跑分类项目就用的是类似结构的数据集踩过的坑和调参经验下文都会讲到。2. 数据集结构与清洗先把文件目录变成可训练的清单拿到数据集第一件事不是急着训练而是确认目录组织方式。10 个做图像识别的项目里至少有 2 个会在数据加载阶段翻车原因都是没看清目录结构。这个 11 类水果数据集常见的组织方式是根目录下每个类别一个文件夹文件夹名就是标签里面放对应水果的图片。这种结构对torchvision.datasets.ImageFolder这类 API 非常友好可以直接读。但友好不等于不用管下面两个环节必须先处理。2.1 目录组织与类别平衡先摸清家底再动手先看类别数量和每类图片张数。数据集的类别是固定的 11 类但每类图片数量未必均匀。有些类可能有 500 张有些可能只有 300 张。在动手之前先写一段脚本统计每类图片数量判断是否需要做类别平衡处理。类别不平衡在分类任务里的后果是模型倾向于预测样本多的类别少样本类别准确率会被压得很低。import os data_root fruit_11_dataset for cls in sorted(os.listdir(data_root)): cls_path os.path.join(data_root, cls) if os.path.isdir(cls_path): n len([f for f in os.listdir(cls_path) if f.lower().endswith((.jpg, .jpeg, .png))]) print(f{cls}: {n})这段脚本做的事很简单遍历根目录下每个子文件夹数出图片文件个数。图片格式判断用扩展名过滤避免把隐藏文件或非图片文件算进去。注意里面的大小写处理.jpg和.JPG都要认。统计完心里就有数了。如果发现最大类和最小类差距超过 3 倍后面训练时就要考虑加权采样或者对少样本类别做更强的数据增强。2.2 标签编码与数据集划分train/val/test 一次性拆好图片分类数据集的标签在文件路径里但训练时不能直接用字符串喂给模型要先做标签到数字 ID 的映射。同时数据集不能只分训练集和验证集最好留出一部分完全不参与训练的测试集最后用测试集评估泛化能力。import random import shutil from pathlib import Path src_root Path(fruit_11_dataset) out_root Path(fruit_11_split) split_ratio {train: 0.7, val: 0.15, test: 0.15} random.seed(42) for cls in sorted(p for p in src_root.iterdir() if p.is_dir()): imgs [p for p in cls.iterdir() if p.suffix.lower() in (.jpg, .jpeg, .png)] random.shuffle(imgs) # 按比例切片 n_train int(len(imgs) * split_ratio[train]) n_val int(len(imgs) * split_ratio[val]) splits { train: imgs[:n_train], val: imgs[n_train:n_train n_val], test: imgs[n_train n_val:] } for split_name, img_list in splits.items(): dst_dir out_root / split_name / cls dst_dir.mkdir(parentsTrue, exist_okTrue) for img in img_list: shutil.copy(img, dst_dir / img.name)这段脚本做两件事一是按 7:1.5:1.5 划分数据集二是把图片复制到train/val/test子目录下对应的类别文件夹。random.seed(42)保证每次运行划分结果一致方便复现实验。为什么要先洗牌再切因为原始数据集的目录可能按拍摄时间或采集批次组织前几十张图可能都是同一光照条件不洗牌会导致验证集和训练集分布不一致。如果后面发现某个类别在测试集上表现异常也可以用同样的 seed 重新划分排查。3. 数据增强与加载器让11类样本在CNN里转起来水果分类有个特点同类水果在形状、颜色上差异不大但光照、背景、拍摄角度变化很大。如果模型只见过固定角度的图换一批照片可能就认不出来了。数据增强是解决这个问题的常用手段本质上是把有限的样本做变换制造更多“看起来合理”的新样本。这个环节做得好不好直接影响最终准确率的上限。3.1 增强策略旋转、翻转、色彩抖动怎么搭配对于水果分类旋转和翻转空间上合理因为水果放在桌上拍摄横着竖着都是同一类。色彩抖动也重要——苹果在不同光照下颜色深浅差异明显不能让模型把“红色”学死。但要注意增强强度不能过火比如把图转到接近倒置模型学到的特征就可能偏离真实分布。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), 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_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里训练集用了 4 种增强水平翻转、角度旋转、色彩抖动、归一化。验证集只做缩放和归一化不加随机增强因为验证集的目的不是制造新样本而是衡量模型在当前权重下的真实表现。如果验证集也做随机增强每次评估的结果都会有随机波动很难判断模型到底进步了没有。RandomRotation(degrees15)我一般控制在正负 15 度以内旋转幅度太大会让语义发生偏移尤其是带有果柄的水果旋转太多会丢失方向性信息。3.2 DataLoader 参数batch size、num_workers 和归一化图像数据加载是训练里最容易被忽略的性能瓶颈。很多人 GPU 跑不满先怀疑模型结构其实问题往往出在 DataLoader 上。num_workers决定用几个子进程加载图片batch_size决定每次迭代送入 GPU 的样本数pin_memory决定是否锁页内存加速 CPU 到 GPU 的传输。from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(rootfruit_11_split/train, transformtrain_transform) val_dataset ImageFolder(rootfruit_11_split/val, transformval_transform) train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue ) print(类别映射:, train_dataset.classes) print(类别数:, len(train_dataset.classes))ImageFolder会按文件夹名字母顺序生成类别映射打印classes能看到类似[apple, banana, ...]的顺序这个顺序就是后面模型输出层的 11 个类别顺序推理时要把索引映射回水果名。batch_size32在 ResNet18 上搭配单张消费级显卡通常刚好显存不够就下调到 16一般不建议超过 64因为水果分类数据本身简单大 batch 反而可能让训练后期收敛变慢。num_workers4是从 4 开始试如果机器内存吃紧就降到 2Windows 平台上如果报多进程相关错误需要把训练代码包在if __name__ __main__:里。4. 模型选型与训练配置从 ResNet18 到迁移学习数据准备好了接下来选模型。这个数据集规模不大每类几百张图总共几千张的量级不需要上 ResNet50 甚至更深的网络。网络越深参数量越大在数据量不够的情况下越容易过拟合。我的经验是先用 ResNet18 跑通流程再考虑迁移学习提升精度这样每一步都有明确的对照。4.1 为什么先选 ResNet18小数据集下的性价比ResNet18 是残差网络里最轻的版本之一18 层深度参数量 1100 万左右单张 224x224 图片前向计算在消费级 GPU 上只要几毫秒。相比 VGG16 的 1.3 亿参数ResNet18 参数少了 10 倍以上但在水果分类这种细粒度要求不高的场景精度差距很小。残差结构的优势是解决了深度网络梯度传播问题即使训练数据不多也不容易发生梯度消失。如果用更深的 ResNet50训练时间翻倍精度可能只提升 1-2 个百分点性价比不高。import torch.nn as nn from torchvision import models def build_model(num_classes11, pretrainedTrue): model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return model model build_model(num_classes11, pretrainedTrue)这里把 ResNet18 最后一层全连接替换成输出 11 类的新全连接层前面的卷积层参数保留 ImageNet 预训练权重。这里有个容易忽略的细节model.fc.in_features是 512这是 ResNet18 最后一层卷积输出的特征维度ResNet50 则是 2048直接把维度写死成 512 以后换模型容易踩坑最好通过in_features动态获取。预训练权重的标准做法是加载 ImageNet 上的权重这个 11 类水果数据集和 ImageNet 中的水果图像有分布重叠迁移学习能明显加速收敛。4.2 训练超参数lr、weight decay、epoch 的常见组合选好模型后训练超参数直接影响结果。这里给一组我试过多次的稳妥组合不是最优解但能让你第一次跑就有可见的收敛趋势。优化器用 Adam 或 SGD 都行迁移学习场景下 Adam 更省心学习率不需要精细调节。import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5) num_epochs 20 best_val_acc 0.0 for epoch in range(num_epochs): model.train() run_loss 0.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() run_loss loss.item() * images.size(0) # 验证 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 correct / total scheduler.step() print(fEpoch {epoch1}/{num_epochs} loss{run_loss/len(train_dataset):.4f} val_acc{val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth)这段代码有几个关键点。lr1e-4是迁移学习的常用起点比从头训练的 1e-3 低一个量级因为预训练权重已经在一个较好的位置上学习率太大容易把学好的特征破坏掉。weight_decay1e-4是 L2 正则防止在小数据集上过拟合。StepLR(step_size5, gamma0.5)表示每 5 个 epoch 学习率减半配合 20 个 epoch 的设定后半程收敛更稳定。验证时用model.eval()和torch.no_grad()是有必要做的两件事前者关闭 Dropout 和 BatchNorm 的训练模式后者关闭梯度计算减少显存占用和计算量。保存模型只保存权重不保存整个模型这样结构代码和权重分离后续换网络后加载也更方便。5. 避坑与常见问题训练11类水果时最容易翻车的五个点数据集的训练流程跑通不难难的是模型表现不稳定时你不知道该改哪里。以下这几条是我用类似水果分类数据集时实际踩过的坑按“现象到原因到解决”的顺序列出来每条都对应一次真实的 debug 经历。5.1 现象验证集准确率震荡不收敛训练 loss 在下降但验证集准确率忽高忽低甚至出现 loss 下降但准确率不涨的情况。排查方向是数据增强和验证集评估的随机性。先确认验证集用的是不含随机变换的固定 transform如果验证时也带了RandomRotation等增强每次评估结果都会不同自然震荡。然后检查drop_last参数训练集最后一个 batch 如果只有 1-2 张图梯度更新波动会很明显。解决方法是把drop_lastTrue加上或者在验证时固定每次的随机种子。遇到这类问题我一般先去掉验证集随机增强再看是否收敛。5.2 现象训练准确率高但验证集明显偏低这说明模型过拟合了。在几千张图片的小数据集上ResNet18 理论容量足够把训练样本全部记住。常见误用是训练轮次太多20 轮已经偏多如果第 10 轮验证准确率就开始原地踏步后面几轮纯属浪费时间。解决手段一是把weight_decay从 1e-4 调到 5e-4增强正则约束二是加数据增强的强度在ColorJitter里加大brightness范围三是提前停止训练每轮验证时把最优模型保存下来训练结束后用保存的模型跑测试集不要用最后一个 epoch 的权重。这个数据集上我的做法是每 5 个 epoch 打印一次验证准确率趋势如果连续 3 次不涨就停。5.3 现象类别间互相误判11 类水果里有几类长得确实接近比如青苹果和梨、橙子和橘子。如果混淆集中在这几对上说明模型学到的特征不够区分。解决思路不是换更大的模型而是检查训练集图片本身——是不是这两类的图片色调高度相似如果是可以在数据增强里加入transforms.GaussianBlur或RandomGrayscale强迫模型关注纹理和形状特征而不是单纯的颜色。我在一个类似项目里给所有训练样本加了 50% 概率的灰度化橙子和橘子的错分率明显下降。另外可以检查这两类的训练图片数量是否明显少于其他类少样本类别需要更强的增强补偿。5.4 现象GPU占用低但训练很慢训练时 GPU 利用率不到 60%CPU 满载每轮要等很久。这是数据加载瓶颈的典型特征。num_workers4在 Windows 上如果数据集放在机械硬盘图片读取速度会成为瓶颈。解决步骤先检查是否因为pin_memoryFalse导致传输慢改为pin_memoryTrue再把num_workers从 4 调到 8前提是你的 CPU 核心数够最后考虑把数据集放到 SSD 上。这个 11 类水果数据集图片格式如果是 PNG读图速度会比 JPG 慢转换过一次格式能明显提速。实测中机械硬盘加num_workers2约等于没开多进程换成 SSD 后提升非常明显。5.5 现象推理时单张图预处理不一致训练跑通了测试集准确率也不错但把单独一张图片拿去做推理时结果不对。最常见的原因是预处理管线和训练时不统一。比如训练时用了Resize((224, 224))推理时忘了缩放直接拿原图尺寸输入或者推理时忘了归一化模型输入分布完全偏移。另一个常见问题是图片通道顺序训练时ToTensor()会把 HWC 转到 CHW如果推理时用 OpenCV 读图默认是 BGR 顺序直接送入模型会颜色错乱。我自己习惯把所有预处理统一封装成一个函数训练和推理共用同一份代码而不是在推理脚本里重新写一遍。6. 验证与进阶用混淆矩阵和分层评估把模型做到可用训练跑了 20 轮验证准确率停在 92% 左右接下来不是继续调参而是先搞清楚模型到底错在哪儿。只看整体准确率会掩盖很多问题比如某个类别准确率只有 60%但因为它样本少对整体准确率的拖累不明显。这一步用混淆矩阵看细粒度表现最直接。下面这段代码用 sklearn 生成混淆矩阵配合seaborn画热力图。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: 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(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstest_dataset.classes, yticklabelstest_dataset.classes) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)看混淆矩阵时我关注两个点一是对角线上的数值是不是每类都明显高于非对角线二是哪两类的误判最多。如果青苹果和梨互相误判说明形状特征没有学充分此时可以针对性地增加纹理增强如果香蕉被误判成芒果很可能是颜色上有重叠光照和色调干扰更大。这张图不仅是分析工具也是向别人证明模型结果的依据——比单独报一个准确率数字更有说服力。进阶操作上模型保存后做推理时可以顺手做一次分层验证按类别分别计算准确率找出最差类别然后针对该类别单独加强增强或补充训练样本。这个数据集的 11 类里必然有一两类准确率偏低不用追求 11 类全部超过 95%但要做到每一类准确率都在 85% 以上才算真正可用。从那以后我每次跑完分类训练都会强制走一遍混淆矩阵 分准确率的流程如果哪个类别准确率偏低就先查数据和增强而不是急着换模型。希望帮到你。本文还有配套的精品资源点击获取
简介这是一份面向深度学习图像识别任务的十一种水果分类数据集覆盖苹果、鳄梨、蓝莓、辣椒、樱桃、猕猴桃、芒果、橙子、岩瓜、草莓、小麦等常见类别适合入门图像分类、训练卷积神经网络以及验证不同网络结构的场景。数据严格按类别分文件夹整理训练集约两千五百余张、测试集约六百余张无需额外标注解压即可直接用于当前主流深度学习框架。压缩包约八百五十余兆共两千个文件图像主体以JPEG格式为主另含少量PNG、WebP、BMP图片并附带类别字典和可视化脚本便于查看类别标签分布和抽样检查图像质量。目前已有超过一千六百人次浏览学习下载后可直接以文件夹名称作为标签省去繁琐的数据清洗流程。对于需要高质量水果图像数据、想专注于模型设计而非数据整理的开发者这份资源能有效提升前期准备效率。1. 11类水果分类数据集为什么分类任务比检测更适合入门做深度学习图像识别很多人一上来就奔着 YOLO 系列去想跑目标检测。但如果你是第一次接触视觉项目我更建议从分类任务起步。这个 11 种水果分类数据集就是一个典型的 11 分类场景苹果、香蕉、葡萄、橙子这类常见水果每类图片数量够用背景不算太干净但也没到故意刁难的程度拿来练手 CNN 或者迁移学习正合适。这个数据集的定位很明确不追求检测框只做整图分类。你给它一张水果图片它输出这个水果属于哪一类。相比检测任务分类不需要标注框只需要文件夹级别的标签数据准备成本低训练收敛也快。适合刚接触深度学习的读者跑通全流程也适合想验证 ResNet、EfficientNet 这类模型在小规模数据上表现的老手做基准测试。我自己第一次跑分类项目就用的是类似结构的数据集踩过的坑和调参经验下文都会讲到。2. 数据集结构与清洗先把文件目录变成可训练的清单拿到数据集第一件事不是急着训练而是确认目录组织方式。10 个做图像识别的项目里至少有 2 个会在数据加载阶段翻车原因都是没看清目录结构。这个 11 类水果数据集常见的组织方式是根目录下每个类别一个文件夹文件夹名就是标签里面放对应水果的图片。这种结构对torchvision.datasets.ImageFolder这类 API 非常友好可以直接读。但友好不等于不用管下面两个环节必须先处理。2.1 目录组织与类别平衡先摸清家底再动手先看类别数量和每类图片张数。数据集的类别是固定的 11 类但每类图片数量未必均匀。有些类可能有 500 张有些可能只有 300 张。在动手之前先写一段脚本统计每类图片数量判断是否需要做类别平衡处理。类别不平衡在分类任务里的后果是模型倾向于预测样本多的类别少样本类别准确率会被压得很低。import os data_root fruit_11_dataset for cls in sorted(os.listdir(data_root)): cls_path os.path.join(data_root, cls) if os.path.isdir(cls_path): n len([f for f in os.listdir(cls_path) if f.lower().endswith((.jpg, .jpeg, .png))]) print(f{cls}: {n})这段脚本做的事很简单遍历根目录下每个子文件夹数出图片文件个数。图片格式判断用扩展名过滤避免把隐藏文件或非图片文件算进去。注意里面的大小写处理.jpg和.JPG都要认。统计完心里就有数了。如果发现最大类和最小类差距超过 3 倍后面训练时就要考虑加权采样或者对少样本类别做更强的数据增强。2.2 标签编码与数据集划分train/val/test 一次性拆好图片分类数据集的标签在文件路径里但训练时不能直接用字符串喂给模型要先做标签到数字 ID 的映射。同时数据集不能只分训练集和验证集最好留出一部分完全不参与训练的测试集最后用测试集评估泛化能力。import random import shutil from pathlib import Path src_root Path(fruit_11_dataset) out_root Path(fruit_11_split) split_ratio {train: 0.7, val: 0.15, test: 0.15} random.seed(42) for cls in sorted(p for p in src_root.iterdir() if p.is_dir()): imgs [p for p in cls.iterdir() if p.suffix.lower() in (.jpg, .jpeg, .png)] random.shuffle(imgs) # 按比例切片 n_train int(len(imgs) * split_ratio[train]) n_val int(len(imgs) * split_ratio[val]) splits { train: imgs[:n_train], val: imgs[n_train:n_train n_val], test: imgs[n_train n_val:] } for split_name, img_list in splits.items(): dst_dir out_root / split_name / cls dst_dir.mkdir(parentsTrue, exist_okTrue) for img in img_list: shutil.copy(img, dst_dir / img.name)这段脚本做两件事一是按 7:1.5:1.5 划分数据集二是把图片复制到train/val/test子目录下对应的类别文件夹。random.seed(42)保证每次运行划分结果一致方便复现实验。为什么要先洗牌再切因为原始数据集的目录可能按拍摄时间或采集批次组织前几十张图可能都是同一光照条件不洗牌会导致验证集和训练集分布不一致。如果后面发现某个类别在测试集上表现异常也可以用同样的 seed 重新划分排查。3. 数据增强与加载器让11类样本在CNN里转起来水果分类有个特点同类水果在形状、颜色上差异不大但光照、背景、拍摄角度变化很大。如果模型只见过固定角度的图换一批照片可能就认不出来了。数据增强是解决这个问题的常用手段本质上是把有限的样本做变换制造更多“看起来合理”的新样本。这个环节做得好不好直接影响最终准确率的上限。3.1 增强策略旋转、翻转、色彩抖动怎么搭配对于水果分类旋转和翻转空间上合理因为水果放在桌上拍摄横着竖着都是同一类。色彩抖动也重要——苹果在不同光照下颜色深浅差异明显不能让模型把“红色”学死。但要注意增强强度不能过火比如把图转到接近倒置模型学到的特征就可能偏离真实分布。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), 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_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里训练集用了 4 种增强水平翻转、角度旋转、色彩抖动、归一化。验证集只做缩放和归一化不加随机增强因为验证集的目的不是制造新样本而是衡量模型在当前权重下的真实表现。如果验证集也做随机增强每次评估的结果都会有随机波动很难判断模型到底进步了没有。RandomRotation(degrees15)我一般控制在正负 15 度以内旋转幅度太大会让语义发生偏移尤其是带有果柄的水果旋转太多会丢失方向性信息。3.2 DataLoader 参数batch size、num_workers 和归一化图像数据加载是训练里最容易被忽略的性能瓶颈。很多人 GPU 跑不满先怀疑模型结构其实问题往往出在 DataLoader 上。num_workers决定用几个子进程加载图片batch_size决定每次迭代送入 GPU 的样本数pin_memory决定是否锁页内存加速 CPU 到 GPU 的传输。from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(rootfruit_11_split/train, transformtrain_transform) val_dataset ImageFolder(rootfruit_11_split/val, transformval_transform) train_loader DataLoader( train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue ) val_loader DataLoader( val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue ) print(类别映射:, train_dataset.classes) print(类别数:, len(train_dataset.classes))ImageFolder会按文件夹名字母顺序生成类别映射打印classes能看到类似[apple, banana, ...]的顺序这个顺序就是后面模型输出层的 11 个类别顺序推理时要把索引映射回水果名。batch_size32在 ResNet18 上搭配单张消费级显卡通常刚好显存不够就下调到 16一般不建议超过 64因为水果分类数据本身简单大 batch 反而可能让训练后期收敛变慢。num_workers4是从 4 开始试如果机器内存吃紧就降到 2Windows 平台上如果报多进程相关错误需要把训练代码包在if __name__ __main__:里。4. 模型选型与训练配置从 ResNet18 到迁移学习数据准备好了接下来选模型。这个数据集规模不大每类几百张图总共几千张的量级不需要上 ResNet50 甚至更深的网络。网络越深参数量越大在数据量不够的情况下越容易过拟合。我的经验是先用 ResNet18 跑通流程再考虑迁移学习提升精度这样每一步都有明确的对照。4.1 为什么先选 ResNet18小数据集下的性价比ResNet18 是残差网络里最轻的版本之一18 层深度参数量 1100 万左右单张 224x224 图片前向计算在消费级 GPU 上只要几毫秒。相比 VGG16 的 1.3 亿参数ResNet18 参数少了 10 倍以上但在水果分类这种细粒度要求不高的场景精度差距很小。残差结构的优势是解决了深度网络梯度传播问题即使训练数据不多也不容易发生梯度消失。如果用更深的 ResNet50训练时间翻倍精度可能只提升 1-2 个百分点性价比不高。import torch.nn as nn from torchvision import models def build_model(num_classes11, pretrainedTrue): model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None) in_features model.fc.in_features model.fc nn.Linear(in_features, num_classes) return model model build_model(num_classes11, pretrainedTrue)这里把 ResNet18 最后一层全连接替换成输出 11 类的新全连接层前面的卷积层参数保留 ImageNet 预训练权重。这里有个容易忽略的细节model.fc.in_features是 512这是 ResNet18 最后一层卷积输出的特征维度ResNet50 则是 2048直接把维度写死成 512 以后换模型容易踩坑最好通过in_features动态获取。预训练权重的标准做法是加载 ImageNet 上的权重这个 11 类水果数据集和 ImageNet 中的水果图像有分布重叠迁移学习能明显加速收敛。4.2 训练超参数lr、weight decay、epoch 的常见组合选好模型后训练超参数直接影响结果。这里给一组我试过多次的稳妥组合不是最优解但能让你第一次跑就有可见的收敛趋势。优化器用 Adam 或 SGD 都行迁移学习场景下 Adam 更省心学习率不需要精细调节。import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5) num_epochs 20 best_val_acc 0.0 for epoch in range(num_epochs): model.train() run_loss 0.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() run_loss loss.item() * images.size(0) # 验证 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 correct / total scheduler.step() print(fEpoch {epoch1}/{num_epochs} loss{run_loss/len(train_dataset):.4f} val_acc{val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth)这段代码有几个关键点。lr1e-4是迁移学习的常用起点比从头训练的 1e-3 低一个量级因为预训练权重已经在一个较好的位置上学习率太大容易把学好的特征破坏掉。weight_decay1e-4是 L2 正则防止在小数据集上过拟合。StepLR(step_size5, gamma0.5)表示每 5 个 epoch 学习率减半配合 20 个 epoch 的设定后半程收敛更稳定。验证时用model.eval()和torch.no_grad()是有必要做的两件事前者关闭 Dropout 和 BatchNorm 的训练模式后者关闭梯度计算减少显存占用和计算量。保存模型只保存权重不保存整个模型这样结构代码和权重分离后续换网络后加载也更方便。5. 避坑与常见问题训练11类水果时最容易翻车的五个点数据集的训练流程跑通不难难的是模型表现不稳定时你不知道该改哪里。以下这几条是我用类似水果分类数据集时实际踩过的坑按“现象到原因到解决”的顺序列出来每条都对应一次真实的 debug 经历。5.1 现象验证集准确率震荡不收敛训练 loss 在下降但验证集准确率忽高忽低甚至出现 loss 下降但准确率不涨的情况。排查方向是数据增强和验证集评估的随机性。先确认验证集用的是不含随机变换的固定 transform如果验证时也带了RandomRotation等增强每次评估结果都会不同自然震荡。然后检查drop_last参数训练集最后一个 batch 如果只有 1-2 张图梯度更新波动会很明显。解决方法是把drop_lastTrue加上或者在验证时固定每次的随机种子。遇到这类问题我一般先去掉验证集随机增强再看是否收敛。5.2 现象训练准确率高但验证集明显偏低这说明模型过拟合了。在几千张图片的小数据集上ResNet18 理论容量足够把训练样本全部记住。常见误用是训练轮次太多20 轮已经偏多如果第 10 轮验证准确率就开始原地踏步后面几轮纯属浪费时间。解决手段一是把weight_decay从 1e-4 调到 5e-4增强正则约束二是加数据增强的强度在ColorJitter里加大brightness范围三是提前停止训练每轮验证时把最优模型保存下来训练结束后用保存的模型跑测试集不要用最后一个 epoch 的权重。这个数据集上我的做法是每 5 个 epoch 打印一次验证准确率趋势如果连续 3 次不涨就停。5.3 现象类别间互相误判11 类水果里有几类长得确实接近比如青苹果和梨、橙子和橘子。如果混淆集中在这几对上说明模型学到的特征不够区分。解决思路不是换更大的模型而是检查训练集图片本身——是不是这两类的图片色调高度相似如果是可以在数据增强里加入transforms.GaussianBlur或RandomGrayscale强迫模型关注纹理和形状特征而不是单纯的颜色。我在一个类似项目里给所有训练样本加了 50% 概率的灰度化橙子和橘子的错分率明显下降。另外可以检查这两类的训练图片数量是否明显少于其他类少样本类别需要更强的增强补偿。5.4 现象GPU占用低但训练很慢训练时 GPU 利用率不到 60%CPU 满载每轮要等很久。这是数据加载瓶颈的典型特征。num_workers4在 Windows 上如果数据集放在机械硬盘图片读取速度会成为瓶颈。解决步骤先检查是否因为pin_memoryFalse导致传输慢改为pin_memoryTrue再把num_workers从 4 调到 8前提是你的 CPU 核心数够最后考虑把数据集放到 SSD 上。这个 11 类水果数据集图片格式如果是 PNG读图速度会比 JPG 慢转换过一次格式能明显提速。实测中机械硬盘加num_workers2约等于没开多进程换成 SSD 后提升非常明显。5.5 现象推理时单张图预处理不一致训练跑通了测试集准确率也不错但把单独一张图片拿去做推理时结果不对。最常见的原因是预处理管线和训练时不统一。比如训练时用了Resize((224, 224))推理时忘了缩放直接拿原图尺寸输入或者推理时忘了归一化模型输入分布完全偏移。另一个常见问题是图片通道顺序训练时ToTensor()会把 HWC 转到 CHW如果推理时用 OpenCV 读图默认是 BGR 顺序直接送入模型会颜色错乱。我自己习惯把所有预处理统一封装成一个函数训练和推理共用同一份代码而不是在推理脚本里重新写一遍。6. 验证与进阶用混淆矩阵和分层评估把模型做到可用训练跑了 20 轮验证准确率停在 92% 左右接下来不是继续调参而是先搞清楚模型到底错在哪儿。只看整体准确率会掩盖很多问题比如某个类别准确率只有 60%但因为它样本少对整体准确率的拖累不明显。这一步用混淆矩阵看细粒度表现最直接。下面这段代码用 sklearn 生成混淆矩阵配合seaborn画热力图。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: 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(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelstest_dataset.classes, yticklabelstest_dataset.classes) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)看混淆矩阵时我关注两个点一是对角线上的数值是不是每类都明显高于非对角线二是哪两类的误判最多。如果青苹果和梨互相误判说明形状特征没有学充分此时可以针对性地增加纹理增强如果香蕉被误判成芒果很可能是颜色上有重叠光照和色调干扰更大。这张图不仅是分析工具也是向别人证明模型结果的依据——比单独报一个准确率数字更有说服力。进阶操作上模型保存后做推理时可以顺手做一次分层验证按类别分别计算准确率找出最差类别然后针对该类别单独加强增强或补充训练样本。这个数据集的 11 类里必然有一两类准确率偏低不用追求 11 类全部超过 95%但要做到每一类准确率都在 85% 以上才算真正可用。从那以后我每次跑完分类训练都会强制走一遍混淆矩阵 分准确率的流程如果哪个类别准确率偏低就先查数据和增强而不是急着换模型。希望帮到你。本文还有配套的精品资源点击获取
