ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

YOLOv5分类训练实战:视觉数据集划分、可视化与混淆矩阵优化

YOLOv5分类训练实战:视觉数据集划分、可视化与混淆矩阵优化 简介面向计算机视觉初学者或 YOLOv5 使用者的交通场景图像分类数据包涵盖汽车、行人、自行车、红绿灯、停车标志等七类常见目标数据已按训练集与验证集划分便于直接开展分类任务训练。压缩包共 2000 个文件以 1998 张 JPG 图片为主体另附 1 个 JSON 类别字典文件和 1 个 Python 可视化脚本整体约 148.28 MB结构简洁明了。训练集包含 6900 张图片验证集包含 2900 张图片比例适合作为 YOLOv5 等主流框架的分类数据集可视化脚本可随机抽取 4 张图片生成预览并保存到当前目录帮助使用者快速确认标注类别与图像质量。目前已有 101 人学习该资源省去了自行采集、清洗和划分数据的步骤拿到即可用于模型训练、实验对比或教学演示。1. 这份视觉分类数据集的真正价值不在数量而在划分拿到data目录的那一刻先别急着数图片。6900 张训练图、2900 张验证图7 个类别覆盖bicycle、car、person、truck、trafficlight、stop、limit30这个规模在今天的深度学习环境下不算大但它属于典型的“可部署级”小样本分类数据集类别之间视觉差异明显背景噪声有限训练曲线收敛快非常适合用来验证分类网络的结构改动、写课程设计或跑通 YOLOv5 分类流程。真正让我觉得这份资源“懂行”的地方是它把 train/val 直接按目录划分好还附带了类别字典 JSON 和可视化脚本。这意味着你不需要自己写数据加载器、不用手撸文件夹遍历拿到手就能开始调模型。适合三类人正在写计算机视觉课程作业的学生、想快速验证分类网络改进效果的算法工程师以及需要一份可靠数据来做 pipeline 测试的嵌入式视觉开发者。2. 数据目录结构、类别字典与 train/val 划分逻辑2.1 目录结构与文件组织解压后你会看到data根目录下只有train、val两个子目录。每个子目录里按类别名再建子文件夹图片存放在对应类别名下。这种“文件夹即标签”的组织方式是 torchvisionImageFolder和 YOLOv5 分类数据加载器默认支持的结构不需要额外的标注文件。data/ ├── train/ │ ├── bicycle/ # 例如 1127.jpg 9093.jpg ... │ ├── car/ │ ├── limit30/ │ ├── person/ │ ├── stop/ │ ├── trafficlight/ │ └── truck/ └── val/ └── 同样的 7 个类别子目录直接数文件数train 共 6900 张val 共 2900 张验证集占比约 29.6%。这个比例比常见的 8:2 略高一点但在小数据集上是合理的7 个类别每类平均训练图不到 1000 张验证集留得越多最终评估指标的置信区间就越窄。我习惯在拿到数据集后先做一步“类别分布核对”防止某个类别的训练/验证比例失衡。用一段简单的 Python 脚本就能跑完import os from collections import Counter root data for split in [train, val]: counter Counter() for cls in os.listdir(f{root}/{split}): cls_path f{root}/{split}/{cls} counter[cls] len(os.listdir(cls_path)) total sum(counter.values()) for cls, num in counter.items(): print(f{split}/{cls}: {num}, 占比 {num/total:.2%})这段代码遍历每个 split 下的类别目录统计图片数量。注意os.listdir返回的列表不会过滤隐藏文件如果后续有人往目录里放了Thumbs.db或.DS_Store计数会偏大。更稳妥的做法是用glob.glob(f{cls_path}/*.jpg)只匹配图片扩展名。检查的意义在于如果某个类别的验证集占比明显偏离 30%你在训练后会得到虚高的准确率部署时才发现对特定类别泛化能力不足。2.2 类别字典文件解读项目附带了一个类别字典内容如下{ 0: bicycle, 1: car, 2: limit30, 3: person, 4: stop, 5: trafficlight, 6: truck }这是一个典型的索引到类别名的映射训练脚本里通常会用idx2class json.load(open(classes.json))读入。这里需要注意两点。第一limit30和stop都属于交通标志但它们被拆成了两个独立类别这意味着训练时网络需要区分“限速 30”和“停止标志”这对小分辨率图片来说是一个容易混淆的点。第二person和bicycle分开但现实场景中行人经常推着自行车所以模型有可能把同时包含两者的图像判成person或bicycle而 “漏掉”另一个目标。如果你后续要做检测而不是分类这个字典可以直接映射成 YOLOv5 的data.yaml类别名顺序保持一致即可。2.3 train/val 划分的合理性6900/2900 的划分比例在原始图片随机分配的逻辑下一般不会出现类别严重偏移。但为了保险建议先验证一下每类在 train 和 val 中的占比是否接近整体比例。我自己遇到过最坑的情况是客户发的数据集里 train 的car全是白天图片val 的car全是傍晚图片结果训练时 loss 降得很好val acc 死活上不去。这个数据集的图片命名比如1127.jpg、9093.jpg看起来像是从某个大库里按时间或编号抽出来的随机性有保障但你不能依赖文件名来推断采样策略。最直接的做法是把上面统计脚本跑一遍算一下每个类别在 train 中的占比# 输出示例 train/person: 1320, 占比 19.13% val/person: 380, 占比 13.10%如果发现某一类在 train 中占 20%在 val 中却只有 10%后续训练完成后需要按类别看 precision/recall不要只认整体准确率。3. Python 可视化脚本随机抽样与 Matplotlib 网格展示3.1 脚本核心逻辑解读项目提供了一个visualize.py作用是从数据集中随机取 4 张图片用 Matplotlib 拼成一张画布展示并保存到当前目录。别看功能简单这种可视化脚本在数据集阶段的意义远超“看一眼长什么样”。它能帮你快速判断三件事图片分辨率是否统一、是否存在损坏文件、各类别图片的拍摄视角和光照差异是否在可接受范围内。常见做法的核心逻辑如下import os import random import matplotlib.pyplot as plt from PIL import Image def visualize_random(splittrain, n4, save_pathvisualize.png): classes [d for d in os.listdir(fdata/{split}) if not d.startswith(.)] selected [] for cls in random.sample(classes, min(n, len(classes))): cls_dir fdata/{split}/{cls} img_name random.choice(os.listdir(cls_dir)) selected.append((cls, os.path.join(cls_dir, img_name))) fig, axes plt.subplots(1, n, figsize(4*n, 4)) for ax, (cls, path) in zip(axes, selected): img Image.open(path).convert(RGB) ax.imshow(img) ax.set_title(f{cls}\n{os.path.basename(path)}, fontsize10) ax.axis(off) plt.tight_layout() plt.savefig(save_path, dpi150) print(fsaved to {save_path})脚本先列出类别文件夹然后用random.sample保证取出的类别不重复接着在每个选中的类别目录里随机抽一张图。这样你看到的 4 张图来自 4 个不同类别能同时观察类别间差异。plt.savefig的dpi150会让保存的 PNG 更清晰方便放进论文或课程报告。需要留意的是random.choice取文件时不会过滤非图片文件建议改成random.choice([f for f in os.listdir(cls_dir) if f.endswith(.jpg)])。3.2 关键参数与运行方式这个脚本只有两个参数值得调整split和n。把splittrain改成val就能检查验证集图片n控制抽样数量。python visualize.py --split val --n 6如果原作者没做命令行参数解析直接在文件尾部改函数参数也能运行。注意保存路径默认是当前目录不要让脚本在只读目录下运行否则会报PermissionError。3.3 可视化维度的三个进阶玩法一次性随机抽 4 张图只能看看表面。我更倾向于对可视化脚本做两个小改进让它变成数据集质量检测工具。第一把抽样改成按类别固定数量比如每个类别抽 2 张7 类就是 14 张用plt.subplots(2, 7)排列一眼看出各类别之间是否有重叠。第二在图片下方打印图片尺寸和通道数如果发现个别图是单通道灰度图Image.open(f).convert(RGB)会替你补成三通道但这治标不治本训练前最好统一转好。第三把损坏图片检测也揉进去用try: img.load() except: print(fbroken: {path})捕捉解压时损坏的 JPEG。这三个改进加起来不到 20 行代码却能避免你训练到一半才被坏图片中断。4. 将数据集投入 YOLOv5 分类训练格式转换、指令与参数坑4.1 数据集与 YOLOv5 分类期望格式的匹配YOLOv5 分类训练支持两种数据格式。第一种是ImageFolder风格也就是当前数据的组织方式直接用--data data/就能读第二种是文本文件每行一张图片路径和类别序号。这份数据集天然符合第一种所以不需要额外转换。目录结构上唯一要确认的是YOLOv5 仓库里的classify/目录下train.py会读取--data指向的根目录然后自动找train和val子目录。如果你的类别目录和这里的完全一致直接把data文件夹放到 YOLOv5 根目录下即可。4.2 启动训练的关键指令在 YOLOv5 仓库根目录执行以下命令python classify/train.py \ --model yolov5s-cls.pt \ --data ../data \ --epochs 50 \ --img 224 \ --batch 32 \ --project runs/classify \ --name car_scene_base几个参数值得单独说明。--model指定预训练模型权重yolov5s-cls.pt是在 ImageNet 上预训练过的分类权重迁移学习起点比随机初始化快很多。--img 224是标准分类输入尺寸这份数据集图片如果原图是 1080P直接缩放会造成信息损失加速不明显改用--img 320是针对性分类精度和速度的常见折中。--batch 32在 8G 显存的卡上配合yolov5s-cls没有问题如果显存不够先降到 16不要动模型宽度。训练过程中最值得观察的不是训练集 loss而是验证集 top1 acc 和 top5 acc。每轮结束时终端会打印类似acc10.921 acc50.987的输出当 top1 acc 不再上升而 top5 acc 还在缓慢爬升时说明网络在部分类别上存在“第二正确答案”的混淆比如把stop判断成limit30这比纯看 loss 曲线更有指导意义。4.3 类别不平衡与边界案例我用这份数据训练时会额外关注person类。行人图在公开数据集中往往有大量半身、背影和遮挡案例而bicycle类则有大量侧面完整自行车这两类的判别边界其实很容易被网络学偏。解决方法也不复杂训练前统计每类图片数量如果某个类别远小于平均值比如低于 700 张在--hyp超参数文件中调整mosaic增强的启用概率或者对较少类别做离线复制。但我不建议直接复制图片因为这会加大过拟合风险。另一个常见的坑是stop和limit30这类交通标志类别如果原图中标志占比较小缩放到 224 之后细节基本丢失。训练时--hsv_v 0.1这类颜色增强参数对红色边框的stop标志影响较大尽量不要加大增强强度。如果验证准确率卡在 95% 不动先检查是不是两类标志被系统性混淆。4.4 从训练日志里定位“难分对”训练结束后最好用验证集做一次完整的评估。YOLOv5 的classify/val.py会输出每个类别的 precision、recall 和更多指标python classify/val.py \ --weights runs/classify/car_scene_base/weights/best.pt \ --data ../data输出的表格里precision低的类别表示模型容易把别的类别误判成它recall低的类别表示它本身容易被漏判。我遇到过truck的 recall 明显低于car的情况原因就是数据集中卡车图片往往有集装箱或货物遮挡和公交车外观相似。针对这类问题可以单独把该类别所有误判图片挑出来看看它们到底被分到了哪个类再决定是补数据还是加一个“truckcar”的合并类。5. 进阶验证用混淆矩阵和错误样本反向改进分类模型最后一章分享一个很实用的小技巧用混淆矩阵 错误样本自动定位模型系统性缺陷比单看准确率高效得多。先写一段脚本加载训练好的模型遍历验证集生成 7×7 混淆矩阵并保存import torch, os, json import numpy as np from PIL import Image from torchvision import transforms model torch.hub.load(ultralytics/yolov5, custom, pathruns/classify/car_scene_base/weights/best.pt, force_reloadTrue) model.eval() tfms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) classes list(json.load(open(classes.json)).values()) conf_mat np.zeros((len(classes), len(classes)), dtypeint) data_root data/val with torch.no_grad(): for cls_id, cls in enumerate(classes): cls_dir f{data_root}/{cls} for fname in os.listdir(cls_dir): img Image.open(f{cls_dir}/{fname}).convert(RGB) inp tfms(img).unsqueeze(0) pred model(inp).argmax(dim1).item() conf_mat[cls_id][pred] 1 np.save(conf_mat.npy, conf_mat)这段代码用torch.hub.load加载自定义模型权重然后手动进行Resize ToTensor Normalize的预处理。关键在预处理必须与训练时一致否则输入分布偏移会显著降低预测置信度。model(inp)前向传播返回的是形状为[1, 7]的逻辑值.argmax(dim1).item()取出预测类别索引。生成的混淆矩阵可以直接用plt.matshow(conf_mat)画热力图也可以打印成表格。拿到混淆矩阵后重点看主对角线之外数值最大的位置。比如person和bicycle互相混淆多说明模型对“人骑车”的组合场景判别能力弱此时可以针对性收集带这种语义的图片。如果stop被分到limit30多优先检查是不是图像分辨率不足导致文字细节丢失。之后再跑一遍下面的代码把每类错误样本的前 10 张复制到指定目录供人工抽查import shutil err_dir errors os.makedirs(err_dir, exist_okTrue) count {c: 0 for c in classes} for true_idx, true_cls in enumerate(classes): cls_dir f{data_root}/{true_cls} for fname in os.listdir(cls_dir): if count[true_cls] 10: break img Image.open(f{cls_dir}/{fname}).convert(RGB) inp tfms(img).unsqueeze(0) pred model(inp).argmax(dim1).item() if pred ! true_idx: shutil.copy(f{cls_dir}/{fname}, f{err_dir}/{true_cls}_pred_{classes[pred]}_{fname}) count[true_cls] 1随后逐个查看错误样本的命名stop_pred_limit30_xxx.jpg这种命名会直接告诉你哪两类在打架。纠错时注意一个原则不要因为某个类别错误率低就忽略它因为 7 类数据量小错误率低也可能只是验证集里该类样本少。把清理过的错误样本补充到训练集之后重新用同样的超参数训练一轮对比 top1 acc 的提升幅度就能量化数据清洗对模型的影响。整个过程十分钟内可以完成但它比盲目调学习率有效得多。本文还有配套的精品资源点击获取
返回列表