图像分类数据集加载与训练实践:从ImageFolder到yolov5/yolov8
简介这是一份面向图像分类入门与目标检测模型训练的图像分类数据集聚焦20种常见小猫品种适合需要快速构建图像分类任务的开发者、学生或AI爱好者。数据已按训练集与测试集划分train目录包含5451张图片test目录包含1352张图片结构清晰可直接配合PyTorch的ImageFolder加载也可作为YOLOv5分类任务的训练数据省去繁琐的整理与标注环节。资源包共2000个文件以1998张jpg图片为主体另附1个Python可视化脚本和1个中文标签JSON字典脚本无需修改即可运行随机传入一张图片即可预览分类效果。压缩包整体大小为224.73MB内容组织规范适合初学者快速上手实践。目前已有227人学习下载对于需要现成、带中文标签的小型图像分类数据集的使用者是不错的选择。1. 图像分类入门先过“数据装载”关20类猫数据集的ImageFolder兼容性图像分类最容易被低估的环节从来不是网络结构而是数据装载。手上有标签但目录混乱或者划分没做就直接训练模型指标再漂亮也无法复现。这套20常见小猫种图像分类数据集把橘猫、布偶、波斯猫等20类猫咪按train/test两个子目录划分好训练集5451张、测试集1352张总大小224MB目录结构与ImageFolder精确吻合。对新手来说它能把精力集中在训练流程而不是洗数据对老手来说自带中文标签字典和可视化脚本适合快速验证图像分类算法、跑迁移学习基线。真正的价值在于数据划分已完成偏差可以追溯。2. 目录布局与ImageFolder加载20类中文标签从文件名到训练索引2.1 数据集的物理结构与ImageFolder的标准约定先看这套数据的目录形态它决定了后续所有加载方式是否成立data/ ├── train/ │ ├── 橘猫/ │ │ ├── Baidu_0005.jpg │ │ ├── Baidu_0096.jpg │ │ └── ... │ ├── 布偶/ │ │ └── ... │ ├── 波斯猫/ │ │ └── ... │ └── # 其余17类 ├── test/ │ ├── 橘猫/ │ │ └── ... │ └── ... └── 猫咪类别.jsontorchvision 的datasets.ImageFolder对目录有一个硬性约定根目录下每个子目录是一个类别子目录名就是标签图片必须直接放在类别目录下不允许再嵌套层级。这套数据是train/test分开的完整数据集而不是只有训练图再加一个csv标注文件因此直接用ImageFolder(data/train)就能生成类别索引不需要额外写一个自定义Dataset来解析标签。这里有一个容易误判的点ImageFolder的class_to_idx并不是按照json文件生成的而是遍历子目录后按目录名字符串排序生成的。中文目录名的排序依赖操作系统locale同一份数据在不同机器上跑出来的类别顺序可能不一致。于是json里写的中文标签与class_to_idx之间可能产生索引错位这是训练前必须先核对的地方。2.2 标签字典的读取、索引对齐与json文件用途json文件存在的意义是保留“中文类别名”这个人类可读层。模型训练最终只用整数索引但推理时要把整数索引还原成“橘猫”“布偶”这样的中文名就靠这个字典。读取时建议写一段最小验证代码把json映射和ImageFolder生成的真实映射打印对照import json from torchvision import datasets, transforms # 1) 读取 json 中文标签字典 with open(猫咪类别.json, r, encodingutf-8) as f: label_map json.load(f) # 兼容两种常见格式{橘猫: 0, ...} 或 {0: 橘猫, ...} if isinstance(list(label_map.values())[0], str): class_names list(label_map.values()) # 数字是 key else: class_names list(label_map.keys()) # 数字是 value # 2) ImageFolder 自己生成的映射 train_dataset datasets.ImageFolder( rootdata/train, transformtransforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), ]) ) print(ImageFolder classes:, train_dataset.classes) print(json 还原 classes:, class_names) print(索引是否对齐:, train_dataset.classes class_names)这段代码的关键在于train_dataset.classes是模型训练时的真实顺序而json只负责显示层转换。如果两行打印结果不一致比如json的顺序是“橘猫、布偶、波斯猫”而ImageFolder按Unicode排序后变成“布偶、波斯猫、橘猫”那么训练得到的结果与中文标签就会整体错一位。我一般会以train_dataset.classes为准重新生成一份group id映射而不是反过来修改目录名。还有一个细节是ImageFolder默认的transform不会做归一化。如果为了训练加了Normalize推理单张图时也要用完全相同的均值和标准差否则输入分布偏移会让softmax输出变得没有区分度。2.3 样本数量与类别数的直观统计拿到数据先不要急着训练先把样本数统计出来。下面这段脚本不读图片内容只遍历文件名因此很快import os def count_files(root): stats {} for cls in os.listdir(root): cls_path os.path.join(root, cls) if os.path.isdir(cls_path): stats[cls] len(os.listdir(cls_path)) return stats train_stats count_files(data/train) test_stats count_files(data/test) print(训练集类别数:, len(train_stats)) print(测试集类别数:, len(test_stats)) print(训练集总数:, sum(train_stats.values())) print(测试集总数:, sum(test_stats.values())) print(平均每类训练样本:, sum(train_stats.values()) // len(train_stats))对照摘要给出的统计训练集5451张、测试集1352张20个类别平均每类约272张训练图。这个数量级对常见品种分类完全是够用的。真正要留意的是类别间数量方差比如某类只有几十张而另一类有四百张这会让模型偏向高频类后续要针对性做类平衡处理。统计输出可以用一张表归档统计项来源预期结果训练集样本总数data/train5451测试集样本总数data/test1352类别数train与test目录20平均每类训练样本数train总数/20约272这个步骤看起来不起眼但它决定了后续是直接训练还是先做重采样。把统计结果存下来之后每次跑实验都能对照查看数据版本训练日志里也更容易排查是谁出了问题。3. 数据划分质量检验训练/测试样本分布与类不平衡核对3.1 划分比例与同分布假设是否成立train与test的比例是5451比1352约为80.2%对19.8%落在图像分类数据集常见的80:20划分区间内。这个比例本身没有明显问题但比例正确不等于“划分质量好”。图像分类里更关键的是同分布假设train和test应该来自同一个原始分布并且图片不重复、近似镜像也不应该跨划分出现。如果原数据是爬虫抓取的同一个来源可能被保存成不同文件名肉眼看不出来但图像内容高度相似。这种重复会让测试集失效因为模型的“记忆”而不是“泛化”被评估了。轻量级做法是用文件MD5做一次全量过滤重一点的可以抽取特征做相似度去重但224MB规模用MD5足够。3.2 用MD5检测跨划分重复图片import os, hashlib def file_md5(path): h hashlib.md5() with open(path, rb) as f: for chunk in iter(lambda: f.read(8192), b): h.update(chunk) return h.hexdigest() def index_by_md5(root): mapping {} for cls in os.listdir(root): cls_path os.path.join(root, cls) if not os.path.isdir(cls_path): continue for fn in os.listdir(cls_path): full_path os.path.join(cls_path, fn) md5 file_md5(full_path) if md5 in mapping: # 同一目录内也可能有重复 mapping[md5].append(full_path) else: mapping[md5] [full_path] return mapping train_index index_by_md5(data/train) test_index index_by_md5(data/test) train_md5s set(train_index.keys()) test_md5s set(test_index.keys()) duplicate_md5s train_md5s test_md5s print(跨 train/test 重复图片组数:, len(duplicate_md5s)) for md5 in list(duplicate_md5s)[:3]: print(train:, train_index[md5], test:, test_index[md5])这段代码按内容摘要而不是文件名去重能立刻暴露train与test之间的泄露。hashlib.md5在这里只需要判断文件是否完全相同不涉及安全性因此速度优先。如果发现重复常见做法是保留train里的图片从test中删除对应项因为测试集必须严格独立。需要注意同一个类目录内部也可能出现几乎一样的连拍图片MD5只能发现字节级完全相同的样本做数据增强时可以把连拍图的去重放在训练策略里考虑。3.3 类不平衡对yolov5和yolov8分类训练的影响类不平衡在分类任务里表现得非常直接样本多的类别loss占比大模型倾向于把边界样本预测为高频类。这个20类猫数据集中常见品种如橘猫、布偶的图片数量大概率明显多于冷门品种训练前至少要看一眼每类样本数的直方图stats {} for cls in os.listdir(data/train): cls_path os.path.join(data/train, cls) if os.path.isdir(cls_path): stats[cls] len(os.listdir(cls_path)) for cls, cnt in sorted(stats.items(), keylambda item: item[1]): print(f{cls}: {cnt})如果发现最少的类只有50张而最多类有500张则在训练yolov5或yolov8分类模型时要考虑以下策略策略适用条件说明按类加权采样类别数量差异 5倍PyTorch的WeightedRandomSampler按样本权重重采样复制增强少数类图绝对数 80张对少数类做随机裁剪、翻转、颜色抖动后重复加入类别权重loss类别数量差异中等在CrossEntropyLoss中传入weight惩罚高频类错误mixup / CutMix整体数据量中等对每批数据做混合增强缓解过拟合多数情况下平均每类270多张的数据集不会出现极端不平衡但仍需跑一遍统计再决定。这个数据集切分好之后我通常先直接训练一个baseline再看混淆矩阵里哪些类互相混淆。猫品种之间外形差异往往是毛色和脸型布偶和暹罗这种本身就有相似特征的类即使样本数均衡也可能混淆这时候问题不在数据划分而在类别定义本身。4. 从ImageFolder到yolov5分类训练加载方式、启动参数与可视化脚本4.1 直接用ImageFolder跑图像分类算法的最小代码ImageFolder最省事的地方在于它同时承担了标签生成和图片解码的组织工作。train和test都已经按类分好所以加载只需要指定transformfrom torch.utils.data import DataLoader from torchvision import datasets, transforms transform_train transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) transform_test 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, transformtransform_train) test_dataset datasets.ImageFolder(data/test, transformtransform_test) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse, num_workers4)四个参数值得说明shuffleTrue只在训练集打开测试集保持顺序便于对齐预测结果和真实标签num_workers4在Windows上如果报错就降为0ImageFolder内部的sample顺序是按目录遍历得到的固定顺序单卡训练时不会有问题RandomCrop与Resize连用让模型看到更多局部纹理比单纯把整图缩放到224×224更稳。如果是Vision Transformer这类图像分类模型通常需要更大分辨率比如384×384这时把Resize改成(384, 384)即可其余结构不用动。4.2 改造成yolov5/yolov8分类数据集并启动训练如果想把这份数据直接喂给yolov5分类训练目录已经合规不需要再改。yolov5官方分类脚本要求的数据结构恰好是datasets/cats/ ├── train/ │ ├── orange_cat/ │ └── ragdoll/ └── test/ ├── orange_cat/ └── ragdoll/把data/train和data/test原样复制或软链到datasets/cats/下即可。启动命令python classify/train.py \ --model yolov5s-cls.pt \ --data datasets/cats \ --epochs 50 \ --img 224 \ --batch 32这里--model yolov5s-cls.pt是yolov5s分类预训练权重首次运行会自动下载--data指向包含train和test的根目录yolov5会自动识别目录下的train和test子目录--img 224控制输入分辨率与上面Resize((224, 224))保持一致。训练完成后用classify/val.py --data datasets/cats --weights best.pt验证top-1准确率。需要注意yolov5分类分支对中文目录名的支持在不同版本里有差异。稳妥做法是把中文标签转换成拼音或英文目录名同时保留json做映射。但我不建议直接重命名目录因为一旦目录名变了json里的键也必须同步改两处很容易出现不一致。更安全的做法是写一个软链层把中文目录在另外一个目录中按英文别名链接过去python - EOF import os, json with open(猫咪类别.json, r, encodingutf-8) as f: label_data json.load(f) class_names list(label_data.values()) if isinstance(list(label_data.values())[0], str) else list(label_data.keys()) for i, name in enumerate(class_names): # 软链目录名例如 cat_00 - 橘猫 print(fcat_{i:02d} - {name}) EOF实际运行时会创建datasets/cats/train/cat_00指向data/train/橘猫的符号链接。这样yolov5看到的是英文目录而我们需要的中文标签仍然完整保留在json里最终推理输出时再做一次反向映射就可以了。相比之下yolov8的yolo classify train datadatasets/cats也接受同样的目录结构所以这套划分好的数据可以在yolov5和yolov8之间无缝切换。4.3 可视化脚本“免改直接跑”的实现原理摘要里说可视化脚本无需更改可以直接运行这里拆解一下它的真实逻辑。脚本要解决的三个问题随机选一张图、显示它的中文标签、把结果保存到当前目录。伪代码如下import random, os, json, matplotlib.pyplot as plt from PIL import Image with open(猫咪类别.json, r, encodingutf-8) as f: label_data json.load(f) class_names list(label_data.values()) if isinstance(list(label_data.values())[0], str) else list(label_data.keys()) cls random.choice(class_names) img_name random.choice(os.listdir(os.path.join(data/train, cls))) img Image.open(os.path.join(data/train, cls, img_name)) plt.imshow(img) plt.title(f{cls} / {img_name}) plt.axis(off) plt.savefig(vis_random.jpg, dpi150, bbox_inchestight) print(saved:, img_name)这段脚本不硬编码图片路径和类别名而是每次从json读取20个类别再随机选目录、随机选图片因此无需改任何参数。注意脚本假设json的key或value中有一侧是中文类别名如果json里存的是“类别名到数字索引”的字典就要像我上面那样先做一次方向判断。真正训练前跑一次这个脚本能最快发现图片是否能正常解码、目录是否为空、json是否与目录名一致这三个问题都是初学者最容易踩的坑。5. 实用改法把随机可视化脚本扩展成单张图片预测验证随机可视化只能确认数据完整性而模型训练完更需要的是“给一张新图输出中文类别和置信度”的验证工具。保留json映射把随机逻辑替换成固定文件输入即可import json, torch from PIL import Image from torchvision import transforms with open(猫咪类别.json, r, encodingutf-8) as f: label_data json.load(f) class_names list(label_data.values()) if isinstance(list(label_data.values())[0], str) else list(label_data.keys()) 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(test.jpg).convert(RGB) x transform(img).unsqueeze(0) model torch.load(best.pt, map_locationcpu) model.eval() with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim1)[0] top3 torch.topk(prob, kmin(3, len(class_names))) for i in range(top3.indices.size(0)): idx top3.indices[i].item() print(f{class_names[idx]:6} {prob[idx].item():.4f})关键点有三处best.pt的加载方式取决于你用的是yolov5分类还是普通PyTorch训练yolov5保存的权重包含模型结构可直接推理自己训练的CNN模型建议同时保存state_dict和类别名数组输入图片必须convert(RGB)避免灰度图通道数不匹配topk的k不能超过类别数20。用这个脚本对test目录里的图批量跑一遍就能快速看出哪些猫种之间互相混淆比如布偶和波斯猫同属长毛类如果经常混在一起说明模型学到的毛色特征多过脸部结构特征。更进一步的验证是让测试集每张图都过一遍模型统计混淆矩阵。如果发现某两类错误率明显偏高常见做法不是盲目加数据而是回到标签定义层确认这两类是否长得过于接近以及标注边界是否清晰。数据划分在这个阶段会再次发挥作用正因为train和test已经按类分好你改任何策略后重新训练跑测试集得到的结果都是可比的。这套数据处理方式在所有图像分类场景里都是同一套思路数据集本身的目录边界已经帮你排掉了分类任务里最脏的活。本文还有配套的精品资源点击获取