简介这是一套面向机器学习与数据挖掘初学者的图像分类数据集包含苹果、香蕉、葡萄、橙子和梨五种常见水果图片每张图片均附有明确分类标签适合用于训练图像识别模型、验证分类算法效果以及作为算法对比的基准数据。包内共有1310个文件主体为1306张jpg图片按类别分目录存放另含2个图片列表文件、1个标签映射json和1个python读取脚本整体约14.07MB轻量易用便于快速下载与解析。目前已有3625人学习使用。借助该资源可实践图像标准化、缩放与增强等预处理流程完成训练集、验证集与测试集划分进而搭建CNN等监督学习模型识别水果还可配合混淆矩阵、精确率、召回率等指标评估模型泛化能力为后续计算机视觉任务积累完整实操经验。1. 水果分类数据集一个能直接跑通图像分类全流程的五类图片资源做图像分类项目最怕的不是模型不会写而是数据集结构混乱、标签对不上、图片尺寸五花八门调了半天发现是数据问题。这份水果分类数据集把苹果、香蕉、葡萄、橙子、梨各归各的子目录文件名也带着类别信息解压之后直接能用省去了大量清洗时间。对刚接触机器学习分类任务的工程师来说它比MNIST和CIFAR更贴近真实场景因为图片不是统一规格对想快速验证CNN结构或迁移学习效果的人来说五类目标不大不小既能看出模型差异又不会让训练等太久。2. 数据集结构与加载先摸清目录组织、标签规则与图片底细不管拿到什么数据集我习惯先不急着训练花半小时把结构、命名、图片基本属性摸一遍。这一步省下来的排查时间后面会成倍回报。2.1 目录组织与文件名规则标签藏在路径和文件名里这份数据集的目录层级很直观。解压后能看到五个子目录分别对应一种水果每类下面就是若干jpg文件。文件名是一串数字加.jpg后缀比如205.jpg和131.jpg这种纯数字命名不是那种带中文或特殊符号的文件名对路径解析很友好。这里有一个关键点值得注意文件名里的数字本身没有业务含义不携带品种、拍摄条件等信息。真正的标签信号来自两级位置——第一级是所在子目录名第二级是图片被归入的子目录归属。也就是说如果你想用glob直接扫目录子目录名就是天然标签如果你把图片全部平铺到一个文件夹里那么必须在扫描时把父目录名记录下来否则标签就丢了。我用一份简单的Python脚本把结构扫出来确认图片数量和分布import os from collections import Counter data_root fruits分类数据集 categories [d for d in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, d))] print(类别列表:, categories) counts Counter() for cat in categories: cat_path os.path.join(data_root, cat) files [f for f in os.listdir(cat_path) if f.lower().endswith((.jpg, .jpeg, .png))] counts[cat] len(files) # 随机抽查三张打印文件名和文件大小 sample_files files[:3] for f in sample_files: size os.path.getsize(os.path.join(cat_path, f)) print(f [{cat}] {f}, {size} bytes) print(每类图片数量:, dict(counts)) print(总图片数:, sum(counts.values()))这段代码做的事是列出顶层目录作为类别候选然后统计每个类别下合法图片格式的文件数量同时打印前三个文件的大小。逻辑上要注意一点os.listdir拿到的顺序不确定所以抽查结果每次可能不一样但这不影响判断数据是否完整。参数方面后缀过滤里我同时包含了.jpg、.jpeg和.png因为有些数据集的图片格式并不统一如果只认.jpg遇到个别.png就会被漏掉后面训练时索引对不上报错都找不到地方。从打印结果能看出这个数据集的特点图片数量不算大每类几十到上百张不等总规模属于小样本范畴。这意味着直接训练深层网络很容易过拟合后面必须靠数据增强或迁移学习来撑。2.2 图片尺寸与色彩模式统一性检查不做这步后面必翻车天然采集的图片不可能全部同尺寸。有些是手机拍的高像素大图有些可能是网络爬虫抓的小图甚至还有带透明通道的PNG混在里面。CNN要求输入张量形状固定所以预处理时通常统一缩放但如果原始图片里混入了损坏文件或单通道灰度图训练时一旦遇到就会炸。我一般会做一个批量检查把尺寸分布和通道数统计出来from PIL import Image import os data_root fruits分类数据集 size_counter {} format_counter {} corrupt_files [] for root, dirs, files in os.walk(data_root): for f in files: if f.lower().endswith((.jpg, .jpeg, .png)): path os.path.join(root, f) try: img Image.open(path) size_counter[img.size] size_counter.get(img.size, 0) 1 format_counter[img.mode] format_counter.get(img.mode, 0) 1 img.load() except Exception as e: corrupt_files.append((path, str(e))) print(尺寸分布:, size_counter) print(色彩模式分布:, format_counter) print(损坏文件数:, len(corrupt_files)) if corrupt_files: print(损坏文件示例:, corrupt_files[:5])这段脚本的关键在于img.load()它会把图片数据真正读入内存而不是只读文件头。如果不调用它PIL的open是惰性的损坏文件要到训练时数据加载器真正解码的那一刻才会报错那时候再回头找就很痛苦。尺寸分布输出类似{(224, 224): 20, (480, 360): 58}这种键值对能看到这个数据集里图片来源比较杂有的已经缩放过有的是原始比例。色彩模式分布里RGB是正常的但如果出现L或RGBA就要注意。L是灰度图直接喂给用三个通道初始化卷积核的模型会报通道数不匹配RGBA需要丢弃alpha通道。常见处理是在加载时统一转成RGB后面预处理那一章我会给具体做法。这一步别偷懒训练跑到一半才发现几十张图片有问题重新跑一轮的时间成本远高于现在多写几行脚本。3. 图像预处理与数据增强让模型吃规格统一的图顺便把小样本撑大数据没洗干净模型再高级也白搭。这一章解决两个问题一是把五花八门的原始图片变成模型能吃的统一张量二是通过数据增强把五类图片的数量和多样性撑起来缓解小样本过拟合。3.1 统一的加载与变换管线从RGB转换到归一化的配置细节无论后面用什么框架预处理的思路是统一的读图→转RGB→缩放→转张量→归一化。这里给出PyTorch风格的transform配置因为torchvision的transforms组合起来最直观也最容易照抄改参数。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), # 统一缩放尺寸 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转概率50% transforms.RandomRotation(15), # 随机旋转正负15度 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), # 转成 [0,1] 范围的张量 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]) ])这个transform序列的每一环都有讲究。Resize((224, 224))是直接把图片拉伸到固定尺寸不管原始宽高比做法粗暴但简单如果想保比例可以用Resize(256)然后CenterCrop(224)但在这个数据集上直接拉伸的效果也不差因为水果主体通常占据画面中心较大区域。RandomHorizontalFlip和RandomRotation是空间变换增强对水果这类物体识别很有效——香蕉横着竖着都能认橘子转个角度也一样。旋转角度设15度是经验参数转太猛会出现黑色边角模型反而学到噪音。ColorJitter的亮度、对比度、饱和度扰动模拟不同光照环境这个数据集里部分图片拍摄条件差异大颜色增强能提升泛化。Normalize用的均值标准差是ImageNet数据集的统计值不是这个水果数据集的统计值。这是迁移学习场景下的惯例如果后续用ImageNet预训练权重输入就应该用ImageNet的均值和标准差去归一化如果从头训练也可以用自己数据集的均值方差但差别不大直接用ImageNet参数更省事。3.2 训练集、验证集、测试集划分比例与随机种子都要定死划分数据集是个容易随手做、也容易出问题的事。常见误区是只分训练集和测试集把验证集省了导致调参时反复用测试集试最终评估结果虚高。import torch from torch.utils.data import Dataset, DataLoader from PIL import Image import os class FruitDataset(Dataset): def __init__(self, root_dir, transformNone): self.samples [] self.transform transform self.classes sorted([d for d in os.listdir(root_dir)]) self.class_to_idx {c: i for i, c in enumerate(self.classes)} for cls in self.classes: cls_dir os.path.join(root_dir, cls) for f in os.listdir(cls_dir): if f.lower().endswith((.jpg, .jpeg, .png)): self.samples.append((os.path.join(cls_dir, f), self.class_to_idx[cls])) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label # 按比例切分 from torch.utils.data import random_split full_dataset FruitDataset(fruits分类数据集, transformtrain_transform) train_size int(0.7 * len(full_dataset)) val_size int(0.15 * len(full_dataset)) test_size len(full_dataset) - train_size - val_size train_ds, val_ds, test_ds random_split( full_dataset, [train_size, val_size, test_size], generatortorch.Generator().manual_seed(42) )这个划分策略是训练集70%、验证集15%、测试集15%随机种子固定为42保证每次运行结果可复现。random_split按比例切分索引不破坏类别平衡性——严格说它不会刻意保持每类占比但因为各类图片数量接近随机切分后的比例偏移可接受。有一个细节我在项目里经常遇到验证集应该用与训练集相同的transform还是用更保守的transform验证集不参与训练不需要数据增强来撑多样性反而应该用最能反映真实分布的变换所以用val_transform不做翻转和颜色扰动。测试集同理。如果验证集也做随机增强评估指标的波动会很大明明没改模型准确率却上蹿下跳纯属自找麻烦。train_loader DataLoader(train_ds, batch_size16, shuffleTrue) val_loader DataLoader(val_ds, batch_size16, shuffleFalse) test_loader DataLoader(test_ds, batch_size16, shuffleFalse)batch_size设16是考虑到图片尺寸224x224浮点计算量和显存开销平衡的结果。如果显存不够改成8也能跑但梯度估计的噪声会变大如果显存充裕32会让训练更快但小数据集上大batch容易收敛到平坦的极值点泛化不一定更好。shuffle只有训练集为True验证和测试集不需要打乱顺序因为评估指标不依赖样本顺序。4. CNN模型训练从零搭建到迁移学习的参数调优实战数据准备好之后核心问题变成选什么模型、怎么训练。这一章先给一个能跑通的小型CNN作基线再介绍迁移学习方案两者对比能帮判断这个数据集到底需要多少模型容量。4.1 从零搭建小型CNN结构设计原则与每层参数含义针对这个小规模数据集我倾向于用一个三层卷积加两层全连接的小网络而不是直接上ResNet。原因很简单数据量摆在那里模型参数越多越容易背住训练集。小网络虽然上限低但训练快、好调参作为baseline足够了。import torch.nn as nn import torch.nn.functional as F class SimpleFruitCNN(nn.Module): def __init__(self, num_classes5): super(SimpleFruitCNN, self).__init__() self.conv1 nn.Conv2d(3, 16, kernel_size3, padding1) self.conv2 nn.Conv2d(16, 32, kernel_size3, padding1) self.conv3 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(kernel_size2, stride2) # 224x224 经过三次池化后是 28x28 self.fc1 nn.Linear(64 * 28 * 28, 256) self.fc2 nn.Linear(256, num_classes) def forward(self, x): x self.pool(F.relu(self.conv1(x))) x self.pool(F.relu(self.conv2(x))) x self.pool(F.relu(self.conv3(x))) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x self.fc2(x) return x网络结构核心在于通道数从16到32再到64逐层翻倍同时特征图尺寸从224降到28。通道数翻倍是因为每经过一次池化空间分辨率减半需要在通道维度上保留更多信息。卷积核选3x3加padding1保证特征图尺寸不缩水只在池化层缩小。全连接层256维是个经验值配五分类输出层容量刚好够不浪费。训练时损失函数用交叉熵优化器选Adam这是图像分类最稳的搭配。以下是训练循环的骨架加了验证集评估的逻辑import torch.optim as optim model SimpleFruitCNN(num_classes5) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) for epoch in range(30): model.train() running_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 每个epoch结束后跑一次验证集 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: 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}: loss{running_loss/len(train_loader):.4f}, val_acc{val_acc:.2f}%)训练循环第11行torch.no_grad()必须带上因为验证阶段不需要计算梯度能省显存并加速。model.eval()切换BatchNorm和Dropout的行为模式——虽然这个小网络没有这些层但养成写eval()的习惯后面用迁移学习时不会忘。Adam初始学习率0.001是默认值在这个任务上通常能快速收敛如果loss震荡厉害下一轮可以降到0.0005。训练30个epoch后观察验证集准确率。正常情况下小CNN大概能到85%到92%之间具体取决于数据增强系数和随机种子。如果验证集准确率比训练集低5个百分点以上说明过拟合开始抬头这时候要么减少网络容量要么增强数据增强强度。4.2 迁移学习用ImageNet预训练权重省下几小时训练时间小CNN跑通后我通常会立刻切到迁移学习方案对比一下效果。这里用torchvision里的ResNet18把最后一层全连接换成五分类输出from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 冻结前面所有层只训练最后的全连接层 for param in model.parameters(): param.requires_grad False num_features model.fc.in_features model.fc nn.Linear(num_features, 5)迁移学习的逻辑很直观ResNet18在前几层学到的是边缘、纹理、颜色斑块这些通用视觉特征这些特征对水果识别同样适用不需要重新学。冻结参数的requires_gradFalse能显著减少计算量只训练最后的全连接层相当于让模型在已有视觉能力的基础上做一次快速适配。model.fc.in_features拿到的是原来最后一层的输入维度这里例子里是512。如果觉得只训练最后一层效果不够可以放开部分层的训练用不同的学习率更新。常见做法是ResNet前半段冻结后半段解冻for name, param in model.named_parameters(): if name.startswith(layer4): # 只解冻最后一个残差块 param.requires_grad True optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.0001)注意这里学习率比从头训练小了一倍还多因为预训练权重已经接近最优调太大一步就可能忘掉之前学到的知识这就是所谓的灾难性遗忘。filter确保Adam只更新需要梯度的参数不会因为冻结层而报错。训练脚本可以复用4.1的循环唯一要改的是优化器和epoch数。预训练模型通常10到15个epoch就能收敛因为它的起点已经很接近好的解。对比下来大概率会看到准确率比小CNN高几个点而且训练时间更短这就是迁移学习的价值。5. 避坑图像分类训练中的六个高频踩坑记录这些坑我基本都踩过一轮有些是数据集本身埋的雷有些是流程设计的问题。每条按现象到原因到解决的思路说清楚遇到类似情况可以直接照着排。5.1 训练中偶发图片解码报错程序直接崩现象训练到某个epoch时数据加载器抛异常提示Image file is truncated或OSError: broken data stream程序中断。换个随机种子跑可能又不出错随机性很强。原因数据集里有几张图片文件不完整可能源头的jpg本身损坏了。PIL默认打开时不一定触发解码错误但数据增强或转Tensor时会真正解码一解码就崩。因为shuffle的关系每次训练遇到的顺序不一样所以表现成随机崩溃。解决在数据加载阶段做一次全量校验。提前用PIL循环打开每张图并load()检查到异常就把路径记下来手动删除或替换损坏文件。这不难前文2.2里的检查脚本就能干这件事。从那以后我拿到任何数据集第一件事就是跑一遍损坏文件扫描绝不偷懒。5.2 验证集准确率很高但真实场景识别效果稀烂现象测试集准确率93%看起来不错但把训练集以外的真实照片喂给模型经常分错。原因测试集和训练集来自同一个数据集图片拍摄风格、背景、设备接近甚至可能有重复图片。模型学到的是这个数据集特有的捷径比如看到白色背景就猜是苹果而不是真正识别水果形状和颜色。解决从数据集里抽出一批图片做隔离测试或者在日常使用中额外收集几十张风格不同的图片做盲测。更稳妥的做法是训练时每个epoch记录验证集准确率如果权重在最后阶段还在上升说明还没有收敛到平台期可以考虑继续训练如果验证集准确率持续高但盲测不行大概率是过拟合到了背景特征应该增强数据增强强度特别是随机裁剪和颜色扰动。5.3 调整随机种子后结果不一致怀疑代码写错现象同一份代码、同一个超参数换了随机种子从85%变成79%差得有点多。原因小数据集本身方差就大模型初始化和数据shuffle顺序的改变都会显著影响最终结果。85%和79%不一定是谁错了更多是数据量少模型对初始化敏感。解决固定随机种子然后跑多次取平均。具体做法是torch.manual_seed(42)统一控制PyTorch的随机源Python的random.seed和NumPy的np.random.seed也要设因为数据加载可能用到它们。固定之后每次结果应该完全一致这时再去对比不同的网络结构或增强策略才公平。最后在报告结果时不要只报最好的一次跑三到五次取均值更诚实。5.4 loss一直不下降训练集准确率也上不去现象训练10个epoch后loss还是在2.0附近徘徊训练集准确率只有40%左右跟猜的差不多。原因最常见的是学习率设置不对或者网络结构有bug。Adam平时不太挑学习率但0.01以上就偏高会造成震荡结构bug则多半是view那一步的特征维度算错了或者卷积输出通道数跟全连接层的输入对不上。解决先看loss初值对不对。五分类交叉熵的随机猜测loss是ln(5)≈1.609如果初始loss远高于这个数说明网络的最后一层输出被bias偏置了或数据标签有问题。如果是维度问题可以把模型输出打印出来检查print(model(torch.randn(1, 3, 224, 224)).shape)输出应该是[1, 5]。如果shape不对对照特征图尺寸计算公式逐层推导我习惯在纸上手算一遍比猜可靠得多。5.5 每类图片的验证集数量很少指标波动巨大现象验证集准确率从90%跳到75%再跳回88%明明代码没动过。每次epoch的验证结果像抽奖。原因图片总量几百张切出15%作为验证集每类可能只有十来张。验证集太小一两张预测错就导致几个百分点的波动评估指标根本不稳定。解决要么把验证集比例提升到20%要么在评估时用多个batch累积统计。还有一个实用技巧是保存验证集图片文件名列表固定验证集内容不随随机种子变化这样调参时对比的是同一份数据上的结果可比较性更强。我在划分数据集时会把索引文件以文本形式落盘避免后续重复划分。5.6 训练时显存不够哪怕batch_size只有8也爆显存现象报错CUDA out of memoryGUP风扇狂转进程直接被kill掉。原因224x224的图片喂给ResNet18本来不费显存但如果用了太大batch累积的中间特征图会撑爆。另一个可能是有其他进程占用了显存或者pytorch缓存没有释放干净。解决先查显存占用nvidia-smi看是不是有僵尸进程有的话kill -9 PID清掉。代码层面可以把torch.backends.cudnn.benchmark设为False减少显存碎片或者把图片尺寸降到160x160五分类任务对分辨率的要求没有想象中高。还有一个实用技巧是训练完一个epoch后手动torch.cuda.empty_cache()虽然对占用帮助有限但能防止长时间训练累积的缓存异常。6. 模型评估与落地混淆矩阵看短板单图推理脚本直接复用模型训练完不等于结束。这一章落到两件事用混淆矩阵找出模型具体在哪些类别上容易混淆然后写一个能对单张图片做预测的推理脚本让它从实验环境走到真实调用。6.1 混淆矩阵与分类报告识别苹果和梨的边界问题准确率是个笼统的指标它不告诉你模型是把苹果认成梨还是把葡萄认成香蕉。混淆矩阵能精确到每个类别的误判方向。from sklearn.metrics import confusion_matrix, classification_report import numpy as np all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(混淆矩阵:) print(cm) class_names [apple, banana, grape, orange, pear] print(classification_report(all_labels, all_preds, target_namesclass_names))混淆矩阵的输出是5x5的数组第i行第j列表示真实类别i的样本被预测成类别j的数量。对角线越亮越好非对角线数字说明具体的混淆方向——如果苹果和梨那两格的数字偏高原因是它们颜色、形状确实接近模型在色彩特征不足时容易走错。classification_report给出的precision、recall、F1值比准确率更有参考价值。如果某一类recall低而precision高说明模型在这个类上预测保守宁可少猜也不轻易给标签这个信息能指导该收集哪些新数据。6.2 单图预测脚本带预处理管线的完整推理函数实验做完了要给别人用或者写进服务需要一个不依赖DataLoader的推理函数。它的作用是把一张任意路径的图片走完整个预处理管线输出五个类别的概率。def predict_image(model, image_path, devicecpu): model.eval() 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]) ]) img Image.open(image_path).convert(RGB) img_tensor transform(img).unsqueeze(0) img_tensor img_tensor.to(device) with torch.no_grad(): outputs model(img_tensor) probs torch.softmax(outputs, dim1) top_prob, top_class torch.max(probs, dim1) class_names [apple, banana, grape, orange, pear] return class_names[top_class.item()], top_prob.item() # 使用示例 pred_class, conf predict_image(model, test_apple.jpg, devicecuda) print(f预测类别: {pred_class}, 置信度: {conf:.3f})这个脚本可以直接复制到项目里用。unsqueeze(0)把单张图片的维度从[C, H, W]变成[1, C, H, W]因为模型接受的是带batch维度的四维张量。torch.softmax把logits转成概率分布torch.max返回概率最大的类别和对应置信度。推理时用torch.no_grad()是必须的跟前文验证集逻辑一样。在推理阶段数据增强是完全不需要的任何随机翻转或颜色扰动都会让结果不稳定。如果想把置信度低于0.6的预测标记为不确定可以在函数里加一个阈值判断这在实际场景里能避免很多误判——模型拿不准的时候直接返回unknown比硬猜一个答案安全。我习惯把所有预处理和推理逻辑封装成独立模块训练脚本和推理脚本共用同一个transform定义避免训练时用一组参数、推理时又用另一组参数这种不一致导致的性能波动最隐蔽。从那以后我每次跑数据集都强制走一遍这个闭环先扫结构、再定预处理、跑基线、看混淆矩阵、封装推理接口五个步骤缺一不可。希望这篇拆解能帮你少走几趟弯路顺利把水果分类这个项目跑通、跑稳。本文还有配套的精品资源点击获取 
长河君 · 2026/10/9 19:48:47 · 阅读 1.2k 水果分类数据集图像分类实战：预处理、CNN训练与迁移学习
简介这是一套面向机器学习与数据挖掘初学者的图像分类数据集包含苹果、香蕉、葡萄、橙子和梨五种常见水果图片每张图片均附有明确分类标签适合用于训练图像识别模型、验证分类算法效果以及作为算法对比的基准数据。包内共有1310个文件主体为1306张jpg图片按类别分目录存放另含2个图片列表文件、1个标签映射json和1个python读取脚本整体约14.07MB轻量易用便于快速下载与解析。目前已有3625人学习使用。借助该资源可实践图像标准化、缩放与增强等预处理流程完成训练集、验证集与测试集划分进而搭建CNN等监督学习模型识别水果还可配合混淆矩阵、精确率、召回率等指标评估模型泛化能力为后续计算机视觉任务积累完整实操经验。1. 水果分类数据集一个能直接跑通图像分类全流程的五类图片资源做图像分类项目最怕的不是模型不会写而是数据集结构混乱、标签对不上、图片尺寸五花八门调了半天发现是数据问题。这份水果分类数据集把苹果、香蕉、葡萄、橙子、梨各归各的子目录文件名也带着类别信息解压之后直接能用省去了大量清洗时间。对刚接触机器学习分类任务的工程师来说它比MNIST和CIFAR更贴近真实场景因为图片不是统一规格对想快速验证CNN结构或迁移学习效果的人来说五类目标不大不小既能看出模型差异又不会让训练等太久。2. 数据集结构与加载先摸清目录组织、标签规则与图片底细不管拿到什么数据集我习惯先不急着训练花半小时把结构、命名、图片基本属性摸一遍。这一步省下来的排查时间后面会成倍回报。2.1 目录组织与文件名规则标签藏在路径和文件名里这份数据集的目录层级很直观。解压后能看到五个子目录分别对应一种水果每类下面就是若干jpg文件。文件名是一串数字加.jpg后缀比如205.jpg和131.jpg这种纯数字命名不是那种带中文或特殊符号的文件名对路径解析很友好。这里有一个关键点值得注意文件名里的数字本身没有业务含义不携带品种、拍摄条件等信息。真正的标签信号来自两级位置——第一级是所在子目录名第二级是图片被归入的子目录归属。也就是说如果你想用glob直接扫目录子目录名就是天然标签如果你把图片全部平铺到一个文件夹里那么必须在扫描时把父目录名记录下来否则标签就丢了。我用一份简单的Python脚本把结构扫出来确认图片数量和分布import os from collections import Counter data_root fruits分类数据集 categories [d for d in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, d))] print(类别列表:, categories) counts Counter() for cat in categories: cat_path os.path.join(data_root, cat) files [f for f in os.listdir(cat_path) if f.lower().endswith((.jpg, .jpeg, .png))] counts[cat] len(files) # 随机抽查三张打印文件名和文件大小 sample_files files[:3] for f in sample_files: size os.path.getsize(os.path.join(cat_path, f)) print(f [{cat}] {f}, {size} bytes) print(每类图片数量:, dict(counts)) print(总图片数:, sum(counts.values()))这段代码做的事是列出顶层目录作为类别候选然后统计每个类别下合法图片格式的文件数量同时打印前三个文件的大小。逻辑上要注意一点os.listdir拿到的顺序不确定所以抽查结果每次可能不一样但这不影响判断数据是否完整。参数方面后缀过滤里我同时包含了.jpg、.jpeg和.png因为有些数据集的图片格式并不统一如果只认.jpg遇到个别.png就会被漏掉后面训练时索引对不上报错都找不到地方。从打印结果能看出这个数据集的特点图片数量不算大每类几十到上百张不等总规模属于小样本范畴。这意味着直接训练深层网络很容易过拟合后面必须靠数据增强或迁移学习来撑。2.2 图片尺寸与色彩模式统一性检查不做这步后面必翻车天然采集的图片不可能全部同尺寸。有些是手机拍的高像素大图有些可能是网络爬虫抓的小图甚至还有带透明通道的PNG混在里面。CNN要求输入张量形状固定所以预处理时通常统一缩放但如果原始图片里混入了损坏文件或单通道灰度图训练时一旦遇到就会炸。我一般会做一个批量检查把尺寸分布和通道数统计出来from PIL import Image import os data_root fruits分类数据集 size_counter {} format_counter {} corrupt_files [] for root, dirs, files in os.walk(data_root): for f in files: if f.lower().endswith((.jpg, .jpeg, .png)): path os.path.join(root, f) try: img Image.open(path) size_counter[img.size] size_counter.get(img.size, 0) 1 format_counter[img.mode] format_counter.get(img.mode, 0) 1 img.load() except Exception as e: corrupt_files.append((path, str(e))) print(尺寸分布:, size_counter) print(色彩模式分布:, format_counter) print(损坏文件数:, len(corrupt_files)) if corrupt_files: print(损坏文件示例:, corrupt_files[:5])这段脚本的关键在于img.load()它会把图片数据真正读入内存而不是只读文件头。如果不调用它PIL的open是惰性的损坏文件要到训练时数据加载器真正解码的那一刻才会报错那时候再回头找就很痛苦。尺寸分布输出类似{(224, 224): 20, (480, 360): 58}这种键值对能看到这个数据集里图片来源比较杂有的已经缩放过有的是原始比例。色彩模式分布里RGB是正常的但如果出现L或RGBA就要注意。L是灰度图直接喂给用三个通道初始化卷积核的模型会报通道数不匹配RGBA需要丢弃alpha通道。常见处理是在加载时统一转成RGB后面预处理那一章我会给具体做法。这一步别偷懒训练跑到一半才发现几十张图片有问题重新跑一轮的时间成本远高于现在多写几行脚本。3. 图像预处理与数据增强让模型吃规格统一的图顺便把小样本撑大数据没洗干净模型再高级也白搭。这一章解决两个问题一是把五花八门的原始图片变成模型能吃的统一张量二是通过数据增强把五类图片的数量和多样性撑起来缓解小样本过拟合。3.1 统一的加载与变换管线从RGB转换到归一化的配置细节无论后面用什么框架预处理的思路是统一的读图→转RGB→缩放→转张量→归一化。这里给出PyTorch风格的transform配置因为torchvision的transforms组合起来最直观也最容易照抄改参数。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), # 统一缩放尺寸 transforms.RandomHorizontalFlip(p0.5), # 随机水平翻转概率50% transforms.RandomRotation(15), # 随机旋转正负15度 transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), # 转成 [0,1] 范围的张量 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]) ])这个transform序列的每一环都有讲究。Resize((224, 224))是直接把图片拉伸到固定尺寸不管原始宽高比做法粗暴但简单如果想保比例可以用Resize(256)然后CenterCrop(224)但在这个数据集上直接拉伸的效果也不差因为水果主体通常占据画面中心较大区域。RandomHorizontalFlip和RandomRotation是空间变换增强对水果这类物体识别很有效——香蕉横着竖着都能认橘子转个角度也一样。旋转角度设15度是经验参数转太猛会出现黑色边角模型反而学到噪音。ColorJitter的亮度、对比度、饱和度扰动模拟不同光照环境这个数据集里部分图片拍摄条件差异大颜色增强能提升泛化。Normalize用的均值标准差是ImageNet数据集的统计值不是这个水果数据集的统计值。这是迁移学习场景下的惯例如果后续用ImageNet预训练权重输入就应该用ImageNet的均值和标准差去归一化如果从头训练也可以用自己数据集的均值方差但差别不大直接用ImageNet参数更省事。3.2 训练集、验证集、测试集划分比例与随机种子都要定死划分数据集是个容易随手做、也容易出问题的事。常见误区是只分训练集和测试集把验证集省了导致调参时反复用测试集试最终评估结果虚高。import torch from torch.utils.data import Dataset, DataLoader from PIL import Image import os class FruitDataset(Dataset): def __init__(self, root_dir, transformNone): self.samples [] self.transform transform self.classes sorted([d for d in os.listdir(root_dir)]) self.class_to_idx {c: i for i, c in enumerate(self.classes)} for cls in self.classes: cls_dir os.path.join(root_dir, cls) for f in os.listdir(cls_dir): if f.lower().endswith((.jpg, .jpeg, .png)): self.samples.append((os.path.join(cls_dir, f), self.class_to_idx[cls])) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label # 按比例切分 from torch.utils.data import random_split full_dataset FruitDataset(fruits分类数据集, transformtrain_transform) train_size int(0.7 * len(full_dataset)) val_size int(0.15 * len(full_dataset)) test_size len(full_dataset) - train_size - val_size train_ds, val_ds, test_ds random_split( full_dataset, [train_size, val_size, test_size], generatortorch.Generator().manual_seed(42) )这个划分策略是训练集70%、验证集15%、测试集15%随机种子固定为42保证每次运行结果可复现。random_split按比例切分索引不破坏类别平衡性——严格说它不会刻意保持每类占比但因为各类图片数量接近随机切分后的比例偏移可接受。有一个细节我在项目里经常遇到验证集应该用与训练集相同的transform还是用更保守的transform验证集不参与训练不需要数据增强来撑多样性反而应该用最能反映真实分布的变换所以用val_transform不做翻转和颜色扰动。测试集同理。如果验证集也做随机增强评估指标的波动会很大明明没改模型准确率却上蹿下跳纯属自找麻烦。train_loader DataLoader(train_ds, batch_size16, shuffleTrue) val_loader DataLoader(val_ds, batch_size16, shuffleFalse) test_loader DataLoader(test_ds, batch_size16, shuffleFalse)batch_size设16是考虑到图片尺寸224x224浮点计算量和显存开销平衡的结果。如果显存不够改成8也能跑但梯度估计的噪声会变大如果显存充裕32会让训练更快但小数据集上大batch容易收敛到平坦的极值点泛化不一定更好。shuffle只有训练集为True验证和测试集不需要打乱顺序因为评估指标不依赖样本顺序。4. CNN模型训练从零搭建到迁移学习的参数调优实战数据准备好之后核心问题变成选什么模型、怎么训练。这一章先给一个能跑通的小型CNN作基线再介绍迁移学习方案两者对比能帮判断这个数据集到底需要多少模型容量。4.1 从零搭建小型CNN结构设计原则与每层参数含义针对这个小规模数据集我倾向于用一个三层卷积加两层全连接的小网络而不是直接上ResNet。原因很简单数据量摆在那里模型参数越多越容易背住训练集。小网络虽然上限低但训练快、好调参作为baseline足够了。import torch.nn as nn import torch.nn.functional as F class SimpleFruitCNN(nn.Module): def __init__(self, num_classes5): super(SimpleFruitCNN, self).__init__() self.conv1 nn.Conv2d(3, 16, kernel_size3, padding1) self.conv2 nn.Conv2d(16, 32, kernel_size3, padding1) self.conv3 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(kernel_size2, stride2) # 224x224 经过三次池化后是 28x28 self.fc1 nn.Linear(64 * 28 * 28, 256) self.fc2 nn.Linear(256, num_classes) def forward(self, x): x self.pool(F.relu(self.conv1(x))) x self.pool(F.relu(self.conv2(x))) x self.pool(F.relu(self.conv3(x))) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x self.fc2(x) return x网络结构核心在于通道数从16到32再到64逐层翻倍同时特征图尺寸从224降到28。通道数翻倍是因为每经过一次池化空间分辨率减半需要在通道维度上保留更多信息。卷积核选3x3加padding1保证特征图尺寸不缩水只在池化层缩小。全连接层256维是个经验值配五分类输出层容量刚好够不浪费。训练时损失函数用交叉熵优化器选Adam这是图像分类最稳的搭配。以下是训练循环的骨架加了验证集评估的逻辑import torch.optim as optim model SimpleFruitCNN(num_classes5) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) for epoch in range(30): model.train() running_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 每个epoch结束后跑一次验证集 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: 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}: loss{running_loss/len(train_loader):.4f}, val_acc{val_acc:.2f}%)训练循环第11行torch.no_grad()必须带上因为验证阶段不需要计算梯度能省显存并加速。model.eval()切换BatchNorm和Dropout的行为模式——虽然这个小网络没有这些层但养成写eval()的习惯后面用迁移学习时不会忘。Adam初始学习率0.001是默认值在这个任务上通常能快速收敛如果loss震荡厉害下一轮可以降到0.0005。训练30个epoch后观察验证集准确率。正常情况下小CNN大概能到85%到92%之间具体取决于数据增强系数和随机种子。如果验证集准确率比训练集低5个百分点以上说明过拟合开始抬头这时候要么减少网络容量要么增强数据增强强度。4.2 迁移学习用ImageNet预训练权重省下几小时训练时间小CNN跑通后我通常会立刻切到迁移学习方案对比一下效果。这里用torchvision里的ResNet18把最后一层全连接换成五分类输出from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 冻结前面所有层只训练最后的全连接层 for param in model.parameters(): param.requires_grad False num_features model.fc.in_features model.fc nn.Linear(num_features, 5)迁移学习的逻辑很直观ResNet18在前几层学到的是边缘、纹理、颜色斑块这些通用视觉特征这些特征对水果识别同样适用不需要重新学。冻结参数的requires_gradFalse能显著减少计算量只训练最后的全连接层相当于让模型在已有视觉能力的基础上做一次快速适配。model.fc.in_features拿到的是原来最后一层的输入维度这里例子里是512。如果觉得只训练最后一层效果不够可以放开部分层的训练用不同的学习率更新。常见做法是ResNet前半段冻结后半段解冻for name, param in model.named_parameters(): if name.startswith(layer4): # 只解冻最后一个残差块 param.requires_grad True optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr0.0001)注意这里学习率比从头训练小了一倍还多因为预训练权重已经接近最优调太大一步就可能忘掉之前学到的知识这就是所谓的灾难性遗忘。filter确保Adam只更新需要梯度的参数不会因为冻结层而报错。训练脚本可以复用4.1的循环唯一要改的是优化器和epoch数。预训练模型通常10到15个epoch就能收敛因为它的起点已经很接近好的解。对比下来大概率会看到准确率比小CNN高几个点而且训练时间更短这就是迁移学习的价值。5. 避坑图像分类训练中的六个高频踩坑记录这些坑我基本都踩过一轮有些是数据集本身埋的雷有些是流程设计的问题。每条按现象到原因到解决的思路说清楚遇到类似情况可以直接照着排。5.1 训练中偶发图片解码报错程序直接崩现象训练到某个epoch时数据加载器抛异常提示Image file is truncated或OSError: broken data stream程序中断。换个随机种子跑可能又不出错随机性很强。原因数据集里有几张图片文件不完整可能源头的jpg本身损坏了。PIL默认打开时不一定触发解码错误但数据增强或转Tensor时会真正解码一解码就崩。因为shuffle的关系每次训练遇到的顺序不一样所以表现成随机崩溃。解决在数据加载阶段做一次全量校验。提前用PIL循环打开每张图并load()检查到异常就把路径记下来手动删除或替换损坏文件。这不难前文2.2里的检查脚本就能干这件事。从那以后我拿到任何数据集第一件事就是跑一遍损坏文件扫描绝不偷懒。5.2 验证集准确率很高但真实场景识别效果稀烂现象测试集准确率93%看起来不错但把训练集以外的真实照片喂给模型经常分错。原因测试集和训练集来自同一个数据集图片拍摄风格、背景、设备接近甚至可能有重复图片。模型学到的是这个数据集特有的捷径比如看到白色背景就猜是苹果而不是真正识别水果形状和颜色。解决从数据集里抽出一批图片做隔离测试或者在日常使用中额外收集几十张风格不同的图片做盲测。更稳妥的做法是训练时每个epoch记录验证集准确率如果权重在最后阶段还在上升说明还没有收敛到平台期可以考虑继续训练如果验证集准确率持续高但盲测不行大概率是过拟合到了背景特征应该增强数据增强强度特别是随机裁剪和颜色扰动。5.3 调整随机种子后结果不一致怀疑代码写错现象同一份代码、同一个超参数换了随机种子从85%变成79%差得有点多。原因小数据集本身方差就大模型初始化和数据shuffle顺序的改变都会显著影响最终结果。85%和79%不一定是谁错了更多是数据量少模型对初始化敏感。解决固定随机种子然后跑多次取平均。具体做法是torch.manual_seed(42)统一控制PyTorch的随机源Python的random.seed和NumPy的np.random.seed也要设因为数据加载可能用到它们。固定之后每次结果应该完全一致这时再去对比不同的网络结构或增强策略才公平。最后在报告结果时不要只报最好的一次跑三到五次取均值更诚实。5.4 loss一直不下降训练集准确率也上不去现象训练10个epoch后loss还是在2.0附近徘徊训练集准确率只有40%左右跟猜的差不多。原因最常见的是学习率设置不对或者网络结构有bug。Adam平时不太挑学习率但0.01以上就偏高会造成震荡结构bug则多半是view那一步的特征维度算错了或者卷积输出通道数跟全连接层的输入对不上。解决先看loss初值对不对。五分类交叉熵的随机猜测loss是ln(5)≈1.609如果初始loss远高于这个数说明网络的最后一层输出被bias偏置了或数据标签有问题。如果是维度问题可以把模型输出打印出来检查print(model(torch.randn(1, 3, 224, 224)).shape)输出应该是[1, 5]。如果shape不对对照特征图尺寸计算公式逐层推导我习惯在纸上手算一遍比猜可靠得多。5.5 每类图片的验证集数量很少指标波动巨大现象验证集准确率从90%跳到75%再跳回88%明明代码没动过。每次epoch的验证结果像抽奖。原因图片总量几百张切出15%作为验证集每类可能只有十来张。验证集太小一两张预测错就导致几个百分点的波动评估指标根本不稳定。解决要么把验证集比例提升到20%要么在评估时用多个batch累积统计。还有一个实用技巧是保存验证集图片文件名列表固定验证集内容不随随机种子变化这样调参时对比的是同一份数据上的结果可比较性更强。我在划分数据集时会把索引文件以文本形式落盘避免后续重复划分。5.6 训练时显存不够哪怕batch_size只有8也爆显存现象报错CUDA out of memoryGUP风扇狂转进程直接被kill掉。原因224x224的图片喂给ResNet18本来不费显存但如果用了太大batch累积的中间特征图会撑爆。另一个可能是有其他进程占用了显存或者pytorch缓存没有释放干净。解决先查显存占用nvidia-smi看是不是有僵尸进程有的话kill -9 PID清掉。代码层面可以把torch.backends.cudnn.benchmark设为False减少显存碎片或者把图片尺寸降到160x160五分类任务对分辨率的要求没有想象中高。还有一个实用技巧是训练完一个epoch后手动torch.cuda.empty_cache()虽然对占用帮助有限但能防止长时间训练累积的缓存异常。6. 模型评估与落地混淆矩阵看短板单图推理脚本直接复用模型训练完不等于结束。这一章落到两件事用混淆矩阵找出模型具体在哪些类别上容易混淆然后写一个能对单张图片做预测的推理脚本让它从实验环境走到真实调用。6.1 混淆矩阵与分类报告识别苹果和梨的边界问题准确率是个笼统的指标它不告诉你模型是把苹果认成梨还是把葡萄认成香蕉。混淆矩阵能精确到每个类别的误判方向。from sklearn.metrics import confusion_matrix, classification_report import numpy as np all_preds [] all_labels [] model.eval() with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) print(混淆矩阵:) print(cm) class_names [apple, banana, grape, orange, pear] print(classification_report(all_labels, all_preds, target_namesclass_names))混淆矩阵的输出是5x5的数组第i行第j列表示真实类别i的样本被预测成类别j的数量。对角线越亮越好非对角线数字说明具体的混淆方向——如果苹果和梨那两格的数字偏高原因是它们颜色、形状确实接近模型在色彩特征不足时容易走错。classification_report给出的precision、recall、F1值比准确率更有参考价值。如果某一类recall低而precision高说明模型在这个类上预测保守宁可少猜也不轻易给标签这个信息能指导该收集哪些新数据。6.2 单图预测脚本带预处理管线的完整推理函数实验做完了要给别人用或者写进服务需要一个不依赖DataLoader的推理函数。它的作用是把一张任意路径的图片走完整个预处理管线输出五个类别的概率。def predict_image(model, image_path, devicecpu): model.eval() 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]) ]) img Image.open(image_path).convert(RGB) img_tensor transform(img).unsqueeze(0) img_tensor img_tensor.to(device) with torch.no_grad(): outputs model(img_tensor) probs torch.softmax(outputs, dim1) top_prob, top_class torch.max(probs, dim1) class_names [apple, banana, grape, orange, pear] return class_names[top_class.item()], top_prob.item() # 使用示例 pred_class, conf predict_image(model, test_apple.jpg, devicecuda) print(f预测类别: {pred_class}, 置信度: {conf:.3f})这个脚本可以直接复制到项目里用。unsqueeze(0)把单张图片的维度从[C, H, W]变成[1, C, H, W]因为模型接受的是带batch维度的四维张量。torch.softmax把logits转成概率分布torch.max返回概率最大的类别和对应置信度。推理时用torch.no_grad()是必须的跟前文验证集逻辑一样。在推理阶段数据增强是完全不需要的任何随机翻转或颜色扰动都会让结果不稳定。如果想把置信度低于0.6的预测标记为不确定可以在函数里加一个阈值判断这在实际场景里能避免很多误判——模型拿不准的时候直接返回unknown比硬猜一个答案安全。我习惯把所有预处理和推理逻辑封装成独立模块训练脚本和推理脚本共用同一个transform定义避免训练时用一组参数、推理时又用另一组参数这种不一致导致的性能波动最隐蔽。从那以后我每次跑数据集都强制走一遍这个闭环先扫结构、再定预处理、跑基线、看混淆矩阵、封装推理接口五个步骤缺一不可。希望这篇拆解能帮你少走几趟弯路顺利把水果分类这个项目跑通、跑稳。本文还有配套的精品资源点击获取