ARTICLE DETAIL

资讯详情

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

水果图像分类实战:从数据预处理到CNN迁移学习全流程

水果图像分类实战:从数据预处理到CNN迁移学习全流程 简介这份水果分类数据集rar压缩包面向机器学习、计算机视觉及数据挖掘方向的初学者与开发者可用于图像分类模型的训练与验证。包内共1310个文件包括1306张苹果、香蕉、葡萄、橙子、梨五种水果的jpg图片以及list标签文件、json配置和py脚本方便快速开展数据加载与预处理。资源整体约14.07MB体量适中适合本地实验或教学演示。已有3625人学习下载具有一定参考热度。通过该数据集读者可以实践从图像标准化、数据集划分到CNN等分类模型搭建、评估的完整流程直观理解特征提取与类别判别的关系是一份兼顾理论与动手操作的入门级计算机视觉练习素材。1. 水果分类数据集为什么我建议你从它开始练手做图像分类的读者应该都有体会最难的不是模型结构而是找一个规模合适、标签干净、能快速跑通全流程的数据集。ImageNet和COCO太重CIFAR-10太抽象而这份 fruits分类数据集 刚好卡在“练手最舒服”的档位——苹果、香蕉、葡萄、橙子、梨五类常见水果图片按类别组织拿来做机器学习分类、数据挖掘特征分析、入门卷积神经网络都非常顺。我拆过不少数据集坦率说这个包更适合拿来跑通“数据预处理 → 模型训练 → 评估调参”的整套流程而不是直接上生产。如果你正在做课程设计、刚入门机器学习想找一个能快速复现的分类任务或者想验证一下自己写的图像分类代码有没有问题这份资源值得下载。2. 解压与预处理先把数据结构和标签规则吃透2.1 解压后先做的事统计文件、看清目录、确认标签拿到fruits分类数据集.rar之后第一步不是急着写模型而是先把压缩包里的真实结构摸清楚。这个数据集的图片文件名是纯数字编号比如205.jpg、131.jpg、212.jpg也就是说文件名本身不携带类别信息类别是靠目录结构来区分的。很多新手在这里容易犯一个错——以为文件名里的编号和类别有对应关系翻半天文件找不到规律其实规律在文件夹上。我的习惯是解压后先跑一段脚本把目录结构、每类图片数量、图片尺寸范围一次性统计出来import os from collections import Counter from PIL import Image data_root fruits_dataset # 解压后的根目录 class_counts Counter() size_stats {min_w: 99999, min_h: 99999, max_w: 0, max_h: 0} for class_name in os.listdir(data_root): class_path os.path.join(data_root, class_name) if not os.path.isdir(class_path): continue imgs os.listdir(class_path) class_counts[class_name] len(imgs) for img_name in imgs: img_path os.path.join(class_path, img_name) with Image.open(img_path) as im: w, h im.size size_stats[min_w] min(size_stats[min_w], w) size_stats[min_h] min(size_stats[min_h], h) size_stats[max_w] max(size_stats[max_w], w) size_stats[max_h] max(size_stats[max_h], h) print(类别统计:, dict(class_counts)) print(尺寸范围:, size_stats)这段代码做的事很直接遍历根目录下每个子文件夹把子文件夹名当作类别名统计每类图片数量同时用 PIL 库读取所有图片的尺寸求出宽高的最小值和最大值。参数说明data_root是解压后的数据集路径需要改成你自己的真实路径class_name在这里直接作为标签因为数据集的目录名就是 apple、banana、grape 这类英文类别名。跑完这段脚本你会看到两类关键信息——第一每类图片数量是否均衡第二图片尺寸是否一致。这两点直接影响后面的模型设计。如果发现某类图片特别少那后面训练时就得考虑数据增强来补足如果图片尺寸不统一那就需要统一 resize。实际拆包时我遇到过一个情况有些图片虽然后缀是 .jpg但编码格式异常PIL 打开时直接报错。这种问题洽洽是后续训练时最容易翻车的地方所以统计脚本里通常会加一个 try-except把打不开的图片单独列出来。2.2 图片统一尺寸与归一化两个必须做的预处理水果分类任务的图片来源比较杂有些是手机拍的有些是网络爬的尺寸可能从几百像素到上千像素都有。CNN 的输入层是固定尺寸的所以第一步就是缩放。缩放不是随便拉一下就行要注意保持宽高比避免水果形状被拉伸变形。我的预处理流程分三步先缩放短边到目标尺寸再中心裁剪最后做归一化。下面这段代码用 PyTorch 的torchvision.transforms写import torchvision.transforms as T IMG_SIZE 128 # 统一尺寸可调 train_transform T.Compose([ T.Resize((IMG_SIZE, IMG_SIZE)), # 直接拉伸到固定尺寸 T.RandomHorizontalFlip(p0.5), # 随机水平翻转增强泛化 T.RandomRotation(15), # 随机旋转 ±15 度 T.ColorJitter(brightness0.2, contrast0.2, saturation0.2), # 颜色扰动 T.ToTensor(), # 转 Tensor像素值映射到 [0,1] T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) # ImageNet 标准化 ]) test_transform T.Compose([ T.Resize((IMG_SIZE, IMG_SIZE)), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里说明一下为什么不直接用短边缩放加中心裁剪而是用Resize((128, 128))直接拉伸水果分类这种类别差异明显的任务拉伸对形状特征的破坏有限但代码会简单很多。如果你想更精细可以用T.Resize(128), T.CenterCrop(128)的组合换成短边缩放加裁剪。Normalize用的均值标准差是 ImageNet 的统计值这是迁移学习的惯例因为后面要加载的预训练权重是在 ImageNet 上训的输入分布对齐效果更好。如果你是自己从零训练 CNN也可以用mean0.5, std0.5这种简单的映射方式。2.3 按文件夹生成标签并划分数据集数据准备好之后要把“目录名 → 数字编号”的映射关系固化下来然后按比例切分训练集、验证集、测试集。这个数据集规模不大我习惯按 7:2:1 切测试集单独留出来只在最后评估时用一次。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.class_to_idx {} for idx, class_name in enumerate(sorted(os.listdir(root_dir))): class_path os.path.join(root_dir, class_name) if not os.path.isdir(class_path): continue self.class_to_idx[class_name] idx for img_name in os.listdir(class_path): if not img_name.lower().endswith((.jpg, .jpeg, .png)): continue self.samples.append((os.path.join(class_path, img_name), idx)) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] with Image.open(img_path).convert(RGB) as im: if self.transform: im self.transform(im) return im, label分类标签是通过enumerate(sorted(os.listdir(root_dir)))自动生成的保证 apple 永远是 0、banana 永远是 1不会因为目录遍历顺序不同导致标签错位。注意Image.open(img_path).convert(RGB)这个细节如果数据里有灰度图或者 RGBA 图不转 RGB 后面Normalize会报通道数不匹配的错。这就是我前面说“有些图片格式异常”时最容易踩的地方转 RGB 可以统一通道数。数据划分可以直接用random_split保证每类图片在三个集合里的比例基本一致from torch.utils.data import random_split dataset FruitDataset(root_dirfruits_dataset, transformtrain_transform) train_size int(0.7 * len(dataset)) val_size int(0.2 * 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])这种划分方式有一个隐藏问题random_split是完全随机切分没有按类别分层。如果某类图片本来就少随机切分可能让验证集里这一类一张都没有。更严谨的做法是用StratifiedSampler做分层采样但考虑到这个数据集五类图片数量差别不算悬殊用random_split也够用读者如果在意可以自己改成 sklearn 的train_test_split(stratifyy)来实现分层。3. 两类模型方案从零训练 CNN 还是迁移学习3.1 从零搭建一个小型 CNN结构简单、可控性强数据集规模不大不需要上 ResNet 这种深度网络。我一般会先用一个小型 CNN 跑通流程验证数据预处理和训练管线没问题再换成更大的模型。下面这个结构是专门为 128×128 输入设计的import torch.nn as nn class FruitCNN(nn.Module): def __init__(self, num_classes5): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), # 3x128x128 - 32x128x128 nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 32x64x64 nn.Conv2d(32, 64, kernel_size3, padding1), # 64x64x64 nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 64x32x32 nn.Conv2d(64, 128, kernel_size3, padding1), # 128x32x32 nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 128x16x16 ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d(1), # 128x1x1 nn.Flatten(), nn.Dropout(0.3), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这个网络只有三层卷积加三层池化参数量很小。num_classes5对应五种水果。用AdaptiveAvgPool2d(1)的好处是不管输入图片实际尺寸是多少全连接层的输入维度都是固定的这样即使你中途想换更大尺寸的图片模型结构也不用改。从零训练这个模型在水果分类这种特征差异明显的任务上通常能跑到 85% 左右的准确率。但训练过程比较“玄学”——同样的代码换个学习率可能就从 85% 跌到 60%所以超参数的选择比模型结构更影响最终结果。如果你不想在这个阶段花太多时间调参直接跳到迁移学习是更稳的选择。3.2 迁移学习ResNet18 在新数据集上的威力这个数据集每类图片大概几十张到上百张的量级从零训练 CNN 很容易过拟合。我的建议是换用 ImageNet 预训练的 ResNet18 或 MobileNetV3只替换最后的全连接层微调整个网络。迁移学习在这个任务上几乎是“作弊级”的收益验证集准确率通常能直接到 95% 以上。import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 5) # 替换最后一层 # 解冻全部参数用较小学习率微调 for param in model.parameters(): param.requires_grad True这里我只演示了替换最后一层和解冻全部参数。对于数据量更小的场景一个常见做法是冻结前面所有层只训练最后一层学习率可以设大一点比如 1e-3数据量足够时则建议解冻全部层用更小的学习率比如 1e-4 到 5e-5微调效果更好。# 如果只想训练最后一层改成这样 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True两种做法的取舍冻结前层训练速度快、显存占用小适合快速验证解冻全部层通常准确率更高但训练时间更长而且如果学习率设太大容易把预训练特征破坏掉。我的实践习惯是第一次跑用冻结版拿到一个基线准确率然后解冻全部层用1e-4再微调十几个 epoch通常还能再涨 23 个点。3.3 损失函数与优化器分类任务的标准配置水果分类是典型的多分类任务损失函数用交叉熵优化器用 Adam 或 SGD 都行。Adam 收敛快、不需要太多调参经验适合新手SGD 加动量训练时间更长但最终准确率往往更高。import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30)CrossEntropyLoss内部已经包含了 Softmax 计算所以模型最后一层不需要额外加 Softmax。AdamW比 Adam 多了权重衰减的解耦泛化效果通常更好一点。CosineAnnealingLR让学习率按余弦曲线从 1e-4 降到接近 0在训练后期能帮助 loss 稳定收敛。一个容易忽略的细节weight_decay是 L2 正则化设太大模型会欠拟合设太小起不到约束作用。对这个小数据集我一般先设 1e-4如果发现验证集准确率一直上不去再降到 1e-5 试试。4. 训练、评估与超参数调优用数据说话4.1 训练循环验证集评估与最佳模型保存训练循环没有太多花哨的东西但有一个习惯很关键每个 epoch 结束都要在验证集上跑一次评估并且只保存验证集准确率最高的模型权重。这样即使后面训练过拟合了你还是能拿到一份最佳权重。best_val_acc 0.0 for epoch in range(30): model.train() train_loss, train_correct 0.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() train_loss loss.item() * images.size(0) train_correct (outputs.argmax(1) labels).sum().item() model.eval() val_correct, val_total 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) val_correct (outputs.argmax(1) labels).sum().item() val_total labels.size(0) val_acc val_correct / val_total print(fEpoch {epoch1:02d} | Train Loss {train_loss/len(train_ds):.4f} | Val Acc {val_acc:.4f}) if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_fruit_model.pth)model.train()和model.eval()这两个状态切换容易被新手忽略。Dropout 和 BatchNorm 在训练和推理阶段的行为不同不切换会导致推理结果异常。torch.no_grad()在验证阶段关闭梯度计算省内存也加速。训练时留意训练 loss 和验证准确率的变化曲线。正常情况是训练 loss 稳步下降验证准确率逐步上升如果训练 loss 降到很低但验证准确率不涨反跌那就是过拟合的典型信号。4.2 评估指标准确率只看懂了七成F1 和混淆矩阵才是关键准确率在类别不均衡的数据集上很容易骗人假设苹果图片占 60%即使一个什么都不学的模型全猜苹果准确率也有 60%。所以评估阶段要多看几个指标。from sklearn.metrics import classification_report, confusion_matrix import numpy as np all_preds, all_labels [], [] model.eval() with torch.no_grad(): for images, labels in test_loader: images images.to(device) outputs model(images) all_preds.extend(outputs.argmax(1).cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_names[apple, banana, grape, orange, pear])) conf_mat confusion_matrix(all_labels, all_preds) print(混淆矩阵:\n, conf_mat)classification_report会一次性输出每个类别的 precision、recall、F1 和 support。混淆矩阵则能告诉你哪两类被经常搞混。我拆过不少数据集几乎每个分类项目里都有“混淆对”——比如苹果和梨外形接近橙子和橘子颜色相近混淆矩阵一眼就能看出来。4.3 超参数速查表一份可以直接抄的配置为了节省读者反复试错的时间我把这个数据集上实际跑得比较好的几组配置列在下面。这些参数不是绝对的但作为起点足够稳。参数从零训练 CNN迁移学习 ResNet18图片尺寸128×128128×128 或 224×224优化器AdamWAdamW初始学习率1e-31e-4解冻全层Batch Size1616Epoch5030学习率调度CosineAnnealingLRCosineAnnealingLRWeight Decay1e-41e-4数据增强翻转旋转颜色抖动翻转旋转颜色抖动Batch Size 在这个数据集上不建议设太大因为图片数量本身不多设 32 或 64 会导致每个 epoch 的迭代次数很少模型还没怎么学就进入下一轮了。16 是一个平衡速度与效果的合理值。如果图片不够可以引入 MixUp 或 CutMix 这种高级增强策略但在这个中小规模数据集上收益有限而且会明显增加训练时长。先把基础增强做足效果不够再上高级策略这个优先级顺序是值得记住的。5. 避坑与常见问题五条具体踩坑记录5.1 现象loss 下降缓慢甚至不降原因最常见的是学习率设得太大或太小其次是数据没归一化。如果输入图像的像素值是 0-255 的整数而不是 0-1 的浮点数梯度容易爆炸loss 会剧烈震荡。解决先检查ToTensor()是否在预处理管线里。然后从 1e-3 开始如果 loss 震荡往下降到 1e-4如果 loss 几乎不动往上升到 3e-3。还有一个快速验证方法用一小批数据比如 32 张过拟合看一下——如果 loss 能降到接近 0说明代码没问题是训练配置的问题如果小批数据都降不下去那问题在模型结构或预处理上。5.2 现象训练集准确率很高验证集准确率偏低原因典型的过拟合。数据集每类图片数量有限模型把训练数据的细节特征记住了而不是学到水果的通用特征。解决三步走。第一步增强数据增强强度把随机旋转角度从 15 度加大到 30 度加上随机裁剪第二步加大 Dropout 概率到 0.5第三步提前停止训练或者用我在第 4.1 节写的“保存验证集最优模型”策略不用最后一轮的权重。如果这三个方法用完还过拟合就该考虑换预训练模型了——从零训练的 CNN 在小数据集上过拟合几乎是必然的。5.3 现象某些类别在混淆矩阵中错得离谱原因这些类别视觉特征本来就接近比如苹果和梨都有圆形轮廓、红色或黄色区域或者训练集中这对类别的图片数量差距过大。解决先确认训练集里两类图片数量是否均衡如果不均衡对少的类别做更强的增强或者 Copy-Synth 生成合成样本。如果数量均衡还是分不清可以考虑在最后一层之前的特征向量上做线性判别分析或者换一个表达能力更强的预训练模型。真实场景里碰到这个问题的概率很高不要指望加数据增强就能完全解决——视觉上人眼都分不清的图CNN 也很难突破接受一定的混淆率是务实的做法。5.4 现象加载预训练权重时维度不匹配原因torchvision的ResNet18预训练权重是针对 ImageNet 的 1000 类而替换后的model.fc输出是 5 类加载权重时最后一层尺寸对不上。解决这是预期内的报错不是 bug。正确的加载顺序是先加载预训练模型再替换最后一层。如果先替换最后一层再加载权重就会报错。另外要确认weightsResNet18_Weights.IMAGENET1K_V1这种写法能自动下载权重如果你的网络环境访问不了官方源可以手动下载后通过torch.load(..., map_location...)加载本地权重文件。5.5 现象训练时显存不足原因图片尺寸过大或 batch size 过大。在 128×128 下 ResNet18 的显存占用大概 1-2GB如果机器的显卡只有 2GB 显存很容易被挤爆。解决最直接的办法是 batch size 从 16 降到 8 或 4。如果还不够把图片尺寸从 128 降到 96 或 64准确率会有轻微下降但换来的是能跑通整个流程。还有一个技巧是在验证阶段用torch.no_grad()并且不开梯度能省一部分显存。老实说这个数据集规模不大即使只用 CPU 跑 ResNet18每个 epoch 也就一二十秒没必要非用 GPU 不可。6. 进阶用 t-SNE 验证模型到底学到了什么训练完模型很多人就直接拿测试集算个准确率就结束了。但一线做项目的工程师通常还会多走一步把模型在测试集上提取的特征向量拿出来用 t-SNE 降到二维做可视化。这一步能非常直观地告诉你——五类水果在模型眼里是不是真的分得开。import matplotlib.pyplot as plt from sklearn.manifold import TSNE features, labels [], [] model.eval() with torch.no_grad(): for images, labels_batch in test_loader: images images.to(device) # 取最后一层全连接之前的特征需要修改模型 forward 返回值 feat model.features(images) # CNN 版 feat model.fc(feat) # 如果模型里有 fc features.append(feat.cpu().numpy()) labels.extend(labels_batch.numpy()) features np.concatenate(features, axis0) tsne TSNE(n_components2, random_state42, perplexity10) coords tsne.fit_transform(features) plt.figure(figsize(8, 8)) for i, name in enumerate([apple, banana, grape, orange, pear]): mask np.array(labels) i plt.scatter(coords[mask, 0], coords[mask, 1], s10, labelname) plt.legend() plt.savefig(tsne_fruit_features.png, dpi150)这段代码的核心思路是模型对一张图做的判断不是最后输出层那一下决定的而是前面卷积层提取的抽象特征决定的。把这些特征压到二维空间后如果五类水果各自聚成一团、簇与簇之间有清晰的间隔说明模型学到的特征是有区分度的。如果 t-SNE 图里两类纠缠在一起那么即便整体准确率不错这两类在实际应用时也随时可能翻车。t-SNE 的perplexity参数需要注意它和样本量有关样本量小于 100 时perplexity建议设置在 5-15 之间默认的 30 会导致可视化失真。我见过不少人在小数据集上直接用默认值画出来的图一团乱麻还以为是模型没学好其实调一下perplexity图就分开了。拿到 t-SNE 图之后我通常还会做一次“错误样本检查”把测试集里预测错的图片单独输出成一个目录人眼过一遍看看到底错在哪里。这个习惯救过我很多次——有时候数据标签贴错了有时候图像本身模糊到人眼都很难确认。从那以后我每次拿到新数据集都强制走一遍“统计 → 预处理 → 建模 → 特征可视化 → 错误样本复盘”的流程这五步走完一个分类项目的成色基本就摸透了。希望帮到你。本文还有配套的精品资源点击获取
返回列表