资讯动态

水果图像分类数据集实战:PyTorch+ResNet迁移学习全流程

发布时间:2026/10/19 1:49:00 来源:尧图企业网站定制
简介专为深度学习图像识别与图像分类实践准备的水果图像分类数据集包含苹果、香蕉、樱桃、火龙果、芒果、橘子、菠萝、木瓜共8个常见类别适合入门图像分类、模型训练与迁移学习实验无需额外清洗即可直接投入训练也可用于课程设计或算法验证。数据已按训练集与测试集划分训练集2220张、测试集550张子文件夹以类别命名便于加载和评估同时附带classes字典JSON文件与可视化Python脚本可快速查看分类映射和样本分布降低数据检查成本。整个压缩包共2000个文件以JPEG图像为主1812个jpeg辅以少量webp、png及1个json、1个py脚本体积约636.77MB目录结构清晰适配主流深度学习框架的目录读取方式。目前已有314人学习下载适合作为水果识别、数据增强与分类模型调参的练习数据。1. 水果图像分类数据集8 类、2220 张训练图开箱即用的深度学习数据做深度学习图像识别的人大多有过“找个顺手数据集比调参还难”的经历要么类别太多、图片太杂要么标签格式要自己折腾半天。我拆的这套水果图像分类数据集属于少数“开箱即用”的——8 分类每个类别一个文件夹文件夹名就是类别名附带 classes.json 直接打通整数索引和类名的映射解压后几乎不用任何预处理就能灌进训练脚本。训练集 2220 张、测试集 550 张总大小 644MB苹果、香蕉、樱桃、火龙果、芒果、橘子、菠萝、木瓜八大类覆盖日常水果识别最常见场景。新手能拿它跑通“加载—训练—评估”整条链路老手则可以直接拿它当迁移学习 baseline或者把目录替换成自己的图集做二次开发。2. 目录结构与标签体系从>水果数据集根目录/ ├──>from torchvision import datasets train_dataset datasets.ImageFolder(rootdata-train) test_dataset datasets.ImageFolder(rootdata-test) print(train_dataset.classes) # 8 个类别名按内部排序生成 print(train_dataset.class_to_idx) # 类别名 - 整数索引 print(len(train_dataset), len(test_dataset)) # 2220 550逻辑说明ImageFolder 构造时会自动扫描 root 下的所有子目录把子目录名作为类别名目录内所有图片作为该类的样本。class_to_idx 这个字典就是模型训练时 ground truth 的来源它由 torchvision 内部排序规则生成不一定和 classes.json 里手写的顺序一致。参数说明root 指向>{ 0: 苹果, 1: 香蕉, 2: 樱桃, 3: 火龙果, 4: 芒果, 5: 橘子, 6: 菠萝, 7: 木瓜 }json 格式的好处是任何语言都能直接读不需要额外依赖。它的实际价值在于模型预测输出的是一个整数 index只有通过这个字典才能翻译成人能读的类别名。我一般会在训练脚本最开头就把它读进来顺手建立双向映射import json with open(classes.json, encodingutf-8) as f: classes_map json.load(f) idx_to_label {int(k): v for k, v in classes_map.items()} label_to_idx {v: k for k, v in idx_to_label.items()} print(idx_to_label[4]) # 芒果逻辑说明json.load 读出来的键是字符串类型所以用 int(k) 转成整数键这样才可以直接和模型输出的 tensor index 比较。idx_to_label 给推理阶段显示类名用label_to_idx 给训练阶段构造标签时用两个方向都备着就不用临时查。参数说明classes.json 的路径按实际位置写如果解压后和>import os from collections import Counter def count_images(data_root): counter Counter() for class_name in os.listdir(data_root): class_dir os.path.join(data_root, class_name) if not os.path.isdir(class_dir): continue images [f for f in os.listdir(class_dir) if f.lower().endswith((.jpg, .jpeg))] counter[class_name] len(images) return counter for cls, cnt in sorted(count_images(data-train).items()): print(f{cls}: {cnt})逻辑说明遍历>import matplotlib.pyplot as plt from PIL import Image import random def visualize_grid(data_root, classes, per_class4, seed0): random.seed(seed) fig, axes plt.subplots(len(classes), per_class, figsize(per_class * 2, len(classes) * 2)) for row, cls in enumerate(classes): cls_dir os.path.join(data_root, cls) images [f for f in os.listdir(cls_dir) if f.lower().endswith((.jpg, .jpeg))] sampled random.sample(images, min(per_class, len(images))) for col, img_name in enumerate(sampled): img Image.open(os.path.join(cls_dir, img_name)) axes[row][col].imshow(img) axes[row][col].axis(off) if col 0: axes[row][col].set_ylabel(cls, rotation0, fontsize10) plt.tight_layout() plt.savefig(class_preview_grid.png, dpi150) plt.show() visualize_grid(data-train, [苹果, 香蕉, 樱桃, 火龙果, 芒果, 橘子, 菠萝, 木瓜])逻辑说明外层循环按类别逐行处理内层循环对每个类做随机抽样random.seed(seed) 固定后每次运行抽到同一批图结果可复现。保存出来的 PNG 网格图既方便自己复查也可以直接放进项目文档当数据样例展示。参数说明per_class 控制每类抽几张我会先设 4 看一轮再设 8 细看一遍。random.sample 外面的 min() 保护了图片数不足 4 张的边缘场景不会直接抛异常。3.3 加载器调优ImageFolder 搭配 DataLoader可视化确认数据没问题之后训练脚本就要真正把图片读进来。数据集本身虽然不需要预处理但训练代码通常要补齐 Resize、归一化和 batch batchingfrom torchvision import transforms from torch.utils.data import DataLoader train_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] ), ]) train_dataset datasets.ImageFolder(data-train, transformtrain_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers2, pin_memoryTrue) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers2, pin_memoryTrue)逻辑说明Resize 把所有长宽不一的 JPEG 统一到 224×224满足 ResNet 系列输入要求ToTensor 把 PIL 图像转成 CHW 格式的 float tensor像素值从 0~255 缩放到 0~1Normalize 采用 ImageNet 预训练模型的均值方差这一步在做迁移学习时不能省否则预训练权重面对的数据分布是错的。参数说明batch_size32 在笔记本和常规 GPU 上都能跑num_workers2 适合 4 核左右 CPU 的机器让数据读取和模型计算并行pin_memory 在 GPU 训练时把数据放到锁页内存能减少拷贝耗时。test_loader 的 shuffle 必须保持 False评估时顺序打乱没有任何意义还会干扰结果复现。4. 从 0 到可用准确率ResNet 迁移学习的最小流水线数据准备完毕接下来就是跑训练。2770 张图片的规模从头训练一个深 CNN 几乎必过拟合我的经验是直接用 torchvision 的 ResNet 预训练权重做迁移学习这是最稳、最快出结果的路线也就几行代码的事。4.1 替换分类头8 分类迁移学习的第一步ResNet18 在 ImageNet 上预训练过 1000 类我们要做的第一件事是把最后那个全连接层换成 8 输出。import torch import torch.nn as nn from torchvision import models device torch.device(cuda if torch.cuda.is_available() else cpu) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Linear(in_features, 8) for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True model model.to(device)逻辑说明预训练权重把前面所有卷积层训练成一个通用的特征提取器我们冻结这些层只训练最后新换上的全连接层更新参数的量级从千万降到几千训练速度快得多也不容易在小样本下过拟合。参数说明先通过 model.fc.in_features 取出原始 512 维再赋值给新的 Linear避免硬编码维度如果想换成 resnet50这段代码同样适用。requires_grad 控制反向传播时哪些参数会被更新冻结层会被 autograd 自动忽略。4.2 训练循环交叉熵、Adam 与收敛监控迁移学习场景下优化器只需要接收 fc 层的参数学习率也不需要设太高。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr3e-4) for epoch in range(30): model.train() total_loss, correct, total 0.0, 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() total_loss loss.item() correct (outputs.argmax(dim1) labels).sum().item() total labels.size(0) train_acc correct / total avg_loss total_loss / len(train_loader) print(fepoch {epoch 1:02d} | loss {avg_loss:.4f} | acc {train_acc:.4f})逻辑说明每个 epoch 内迭代全部 batch前向计算得到交叉熵损失反向传播算梯度再更新全连接层参数。outputs.argmax(dim1) 取每个样本预测概率最大的索引和 labels 逐元素比较后统计正确样本数当前 epoch 的 acc 就出来了。参数说明lr3e-4 在只训练最后一层的场景下足够用调大容易震荡epoch 先设 30一般到 15 轮左右 loss 会明显趋平。打印格式里对 epoch 做了两位数补零方便终端日志滚动时快速定位。4.3 用测试集说话评估与权重落盘训练集准确率只能说明模型的记忆能力真正的说服力来自>model.eval() test_correct, test_total 0, 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) test_correct (outputs.argmax(dim1) labels).sum().item() test_total labels.size(0) test_acc test_correct / test_total print(ftest accuracy: {test_acc:.4f}) torch.save(model.state_dict(), fruit8_resnet18.pth)逻辑说明model.eval() 会把 dropout、BN 层切到推理模式再配合 no_grad() 关闭梯度计算评估阶段的显存占用低、速度也快。torch.save 保存的是 state_dict 参数文件不包含优化器和训练状态体积小最适合后续推理。参数说明权重文件后缀用 .pth 或 .pt 都行加载时先用相同的 ResNet18 结构构造模型再 load_state_dict。如果测试集准确率明显低于训练集说明有过拟合或数据泄漏需要回到第 5 章排查。4.4 混淆矩阵看清错误集中在哪个类别只看总准确率很容易被平均数骗过去。水果数据里火龙果和木瓜这类颜色、纹理接近的类别误判率往往不低。所以评估阶段我习惯再补一张混淆矩阵import numpy as np from sklearn.metrics import confusion_matrix, classification_report y_true, y_pred [], [] with torch.no_grad(): for images, labels in test_loader: images images.to(device) outputs model(images) preds outputs.argmax(dim1).cpu().numpy() y_pred.extend(preds) y_true.extend(labels.numpy()) cm confusion_matrix(y_true, y_pred) print(classification_report(y_true, y_pred, target_namestrain_dataset.classes))逻辑说明把整个测试集跑一遍收集真实标签和预测标签然后交给 sklearn 生成混淆矩阵和分类报告。classification_report 会按类输出 precision、recall、f1-score一眼就能看出是哪两个类在互相纠缠。参数说明train_dataset.classes 和 classes.json 的顺序可能不同这里以 ImageFolder 的实际顺序为准因为模型训练时的标签索引就来自它。5. 踩坑记录拿这套 8 分类数据集训练时最容易翻车的 5 个现场这种小规模数据集的问题往往不在模型而在数据本身。下面 5 个坑是我实际跑这套数据时真实遇到或预判会遇到的按“现象 — 原因 — 解决”整理照着排查能省下不少时间。5.1 训练准确率接近满分测试集却掉回七八成现象训练集 acc 一路冲到 0.98测试集只有 0.7 出头dropout 和正则都加过了仍然拉不回来。原因train 和 test 两个文件夹里可能有重复或高度相似的图片。数据采集阶段如果用了连拍或多角度拍摄同一个水果的多张照片很容易被打散到两个目录造成信息泄漏模型在训练阶段就已经“见过”测试图了。解决写一个脚本对 train 和 test 所有图片做 MD5 哈希比对把重复项打印到 csv 里人工确认。我一般会把重复图片从测试集删除或移到训练集一侧让两个集合的边界干净之后重新训练。5.2 classes.json 的顺序和 ImageFolder 的 classes 对不上现象训练完画混淆矩阵发现类名全乱了明明是火龙果的样本预测结果却显示成“1”。原因ImageFolder 内部按 Unicode 排序生成类别索引而 classes.json 的数字是建集时人工指定的顺序两边一旦不一致模型学习的标签就和展示用的类名错位。解决不要依赖 classes.json 的“书写顺序”只信任它的“键值对”。加载完 ImageFolder 后先读 classes.json 得到 idx_to_label再用 train_dataset.classes 做一层对齐验证不一致就以 classes.json 为基准重建标签映射。5.3 解压后 ImageFolder 报 Found 0 images现象代码跑起来直接打印 Found 0 images in subfolders of:>import os import json def generate_classes_json(data_root): categories [d for d in sorted(os.listdir(data_root)) if os.path.isdir(os.path.join(data_root, d))] mapping {str(i): name for i, name in enumerate(categories)} with open(classes_new.json, w, encodingutf-8) as f: json.dump(mapping, f, ensure_asciiFalse, indent2) return mapping generate_classes_json(data-train)逻辑说明用 sorted 排序后生成映射确保 classes_new.json 的顺序和 ImageFolder 内部排序保持一致直接规避掉 5.2 节的标签错位问题。6.2 从冻结最后一层升级到全量微调如果你的新数据集和水果图片差异较大比如换成工业零件表面缺陷或遥感地物只训练最后一层通常不够需要解锁全部层做全局微调for param in model.parameters(): param.requires_grad True optimizer torch.optim.Adam([ {params: model.fc.parameters(), lr: 3e-4}, {params: [p for n, p in model.named_parameters() if fc not in n], lr: 1e-5} ])参数说明这里用了分组学习率分类头保持 3e-4特征提取层降到 1e-5避免大步长破坏预训练特征这是实践里比较稳妥的微调节奏。6.3 数据增强是小数据集最大的后悔药水果数据集每类只有两三百张图我通常会把增强加进训练 transform相当于免费扩充样本量train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2, hue0.05), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])逻辑说明RandomResizedCrop 按随机比例裁剪再缩放到统一尺寸相当于给模型不断提供不同视角的切片ColorJitter 对明暗、饱和度做扰动对水果这类颜色敏感任务的鲁棒性提升很明显。测试集 transform 保持 Resize ToTensor Normalize 即可不要做随机增强。从那以后我每次拿到新的图像数据集都会强制走一遍固定流程先哈希去重再统计类目分布最后随机抽样可视化确认内容和标注一致这三步做完才允许进入训练。这套流程在这份水果数据上帮我省掉了大量返工时间换到其他数据集也同样适用。希望这份拆解能帮到你。本文还有配套的精品资源点击获取

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑