ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

YOLOv5分类实战:5种花数据集训练与ONNX导出全流程

YOLOv5分类实战:5种花数据集训练与ONNX导出全流程 简介这份资源面向计算机视觉入门与进阶开发者提供一套可直接运行的YOLOv5图像分类实战方案用于解决花卉五分类任务中数据准备、模型训练与权重复用的问题。压缩包共约2000个文件整体257.86MB其中1866张jpg图片构成五类花数据集51个yaml与51个py文件分别负责数据配置和模型、训练逻辑另有yml、sh、md、ipynb等辅助脚本与说明文档便于快速复现实验。作者已训练10个epochtop1准确度达到0.91以上并附训练好的权重参数网络尚未完全收敛加大轮次仍有提升空间。使用时只需按目录摆放好datasets数据即可开始训练适合想快速跑通分类流程、验证YOLOv5分类能力或在此基础上做改进实验的读者。目前已有366人学习下载可作为分类任务的入门基线参考。1. YOLOv5 分类实战5 种花分类数据集为什么值得从分类任务切入很多人第一次接触 YOLOv5都是从目标检测开始上来就标框、调 anchor、算 mAP结果环境还没配好就被劝退。但如果你把 YOLOv5 当成一个完整的视觉工具箱来看它其实内置了分类classification能力而且用 5 种花分类数据集来练手是性价比极高的一条路径。原因很直接分类任务不需要标注框数据集现成训练轮次少单卡就能跑能让你把 YOLOv5 的目录结构、数据加载、超参数配置、模型导出这一整套流程走通。等你把分类跑顺了再回头做 yolov5训练自己的数据集 或者 yolov5安全帽数据集 这类检测任务会发现底层逻辑是相通的。这篇文章面向的是想用 YOLOv5 做分类但不知道从哪下手的人也适合已经会检测、想补上分类这一块的工程师。我会把环境配置、数据组织、训练命令、超参数调整、踩坑记录全部拆开讲让你照着就能复现。2. 环境配置与数据准备从 conda yolov5 到 5 种花数据集落盘2.1 为什么用 conda 隔离 YOLOv5 环境YOLOv5 的依赖不算复杂但版本敏感。PyTorch、torchvision、numpy、opencv-python 这几个包版本错一个就可能出现 CUDA 不可用或者训练中途报错。我一般不会在 base 环境里直接装而是单独建一个 conda 环境这样即使后面要跑 yolov5部署 或者树莓派5上部署自己训练的yolov5模型也能快速复制一份干净的环境出来。# 创建名为 yolov5_cls 的 conda 环境指定 python 3.8 conda create -n yolov5_cls python3.8 -y # 激活环境 conda activate yolov5_cls # 安装 PyTorch这里以 CUDA 11.3 为例具体版本按你的显卡驱动来 pip install torch1.12.1cu113 torchvision0.13.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113 # 克隆 YOLOv5 仓库如果你已经有源码包跳过这一步 git clone https://github.com/ultralytics/yolov5.git cd yolov5 # 安装项目依赖 pip install -r requirements.txt逻辑说明先建环境再装 PyTorch是为了避免 conda 自动解析依赖时把 torch 装成 CPU 版本。--extra-index-url指向 PyTorch 官方 wheel 源保证装到的是 CUDA 版本。装完依赖后可以用python -c import torch; print(torch.cuda.is_available())验证 GPU 是否可用返回 True 才算环境没问题。参数说明python3.8是 YOLOv5 比较稳的版本3.9 以上有时会遇到 torchvision 兼容问题。CUDA 版本要和你nvidia-smi里显示的驱动支持版本匹配不是越高越好。2.2 5 种花分类数据集的目录结构怎么摆YOLOv5 的分类任务对数据目录有固定要求不是随便扔几个文件夹就行。它要求每个类别一个子目录训练集和验证集分开结构如下flower_dataset/ ├── train/ │ ├── daisy/ │ ├── dandelion/ │ ├── rose/ │ ├── sunflower/ │ └── tulip/ └── val/ ├── daisy/ ├── dandelion/ ├── rose/ ├── sunflower/ └── tulip/每个类别目录下放对应图片图片格式支持 jpg、png 等常见格式。train 和 val 的比例我一般按 8:2 切如果某类样本特别少至少保证 val 里每类有 10 张以上否则验证准确率波动会很大。import os import random import shutil # 原始数据目录假设所有图片混在一起文件名前缀是类别名 src_dir raw_flowers dst_dir flower_dataset classes [daisy, dandelion, rose, sunflower, tulip] split_ratio 0.8 for cls in classes: # 找到该类别的所有图片 imgs [f for f in os.listdir(src_dir) if f.startswith(cls) and f.endswith(.jpg)] random.shuffle(imgs) n_train int(len(imgs) * split_ratio) train_imgs imgs[:n_train] val_imgs imgs[n_train:] # 创建目录 os.makedirs(os.path.join(dst_dir, train, cls), exist_okTrue) os.makedirs(os.path.join(dst_dir, val, cls), exist_okTrue) # 复制文件 for img in train_imgs: shutil.copy(os.path.join(src_dir, img), os.path.join(dst_dir, train, cls, img)) for img in val_imgs: shutil.copy(os.path.join(src_dir, img), os.path.join(dst_dir, val, cls, img)) print(f{cls}: train{len(train_imgs)}, val{len(val_imgs)})逻辑说明这段脚本假设原始图片文件名以类别名开头比如daisy_001.jpg。如果你的数据来源不同改一下筛选逻辑就行。核心是保证每个类别在 train 和 val 下都有独立目录且目录名就是类别名YOLOv5 会自动读取目录名作为标签。参数说明split_ratio控制训练集比例数据量少的时候可以调到 0.9但 val 不能少于每类 5 张。random.shuffle前建议固定随机种子方便复现。提示如果图片尺寸差异很大YOLOv5 分类训练默认会 resize 到 224x224不需要你手动统一尺寸但极端长宽比的图片建议提前筛掉。3. 训练分类模型yolov5 超参数怎么设、命令怎么敲3.1 分类训练命令与关键参数拆解YOLOv5 的分类训练入口是classify/train.py不是检测的train.py这一点很多人第一次会搞错。基本命令如下python classify/train.py \ --data flower_dataset \ --model yolov5s-cls.pt \ --epochs 50 \ --batch-size 32 \ --imgsz 224 \ --device 0 \ --project runs/train-cls \ --name flower_exp1逻辑说明--data指向数据集根目录里面必须有 train 和 val 两个子目录。--model指定预训练权重yolov5s-cls.pt是分类专用的轻量模型比从零训练收敛快很多。--epochs和--batch-size根据你的显存调整224 输入下 8G 显存跑 32 基本没问题。参数说明--imgsz分类任务常用 224如果你图片细节多可以调到 320 或 448但显存占用会明显上升。--device 0表示用第一块 GPUCPU 训练去掉这个参数或者写--device cpu。--project和--name决定输出目录训练日志、权重、混淆矩阵都会存在里面。训练过程中你会看到类似这样的输出Epoch gpu_mem loss lr acc1 acc5 1/50 2.1G 0.892 0.001 0.612 0.923 10/50 2.1G 0.431 0.001 0.834 0.981 50/50 2.1G 0.187 0.0001 0.921 0.998acc1 是 top-1 准确率acc5 是 top-5 准确率。5 种花分类任务里acc1 到 0.9 以上基本就够用了如果卡在 0.7 左右上不去先检查数据有没有标错再考虑调超参数。3.2 超参数调整学习率、批大小、优化器怎么选YOLOv5 分类的默认超参数在data/hyps/hyp.scratch-low.yaml里但分类任务实际用的是classify/train.py内部的默认值。我一般会关注这几个参数默认值建议调整范围说明lr00.0010.0001 ~ 0.01初始学习率预训练模型用 0.001 比较稳lrf0.010.001 ~ 0.1最终学习率系数控制衰减终点momentum0.90.8 ~ 0.95SGD 动量Adam 不用管weight_decay0.00050.0001 ~ 0.001权重衰减防过拟合warmup_epochs31 ~ 5预热轮次小数据集可以减到 1label_smoothing0.00.0 ~ 0.1标签平滑过拟合时加到 0.1如果你发现训练集准确率很高但验证集上不去先把label_smoothing调到 0.1再把weight_decay加到 0.001。如果训练 loss 下降很慢把lr0提到 0.01 试试但不要超过 0.01否则容易震荡。# 带超参数覆盖的训练命令 python classify/train.py \ --data flower_dataset \ --model yolov5s-cls.pt \ --epochs 80 \ --batch-size 16 \ --imgsz 224 \ --lr0 0.001 \ --lrf 0.01 \ --label-smoothing 0.1 \ --weight-decay 0.001 \ --device 0逻辑说明这里显式覆盖了几个关键超参数--label-smoothing 0.1和--weight-decay 0.001是应对过拟合的组合拳。--batch-size降到 16 是因为加了正则化后梯度噪声变大小 batch 有时反而更稳。参数说明--lr0和--lrf共同决定学习率曲线实际学习率从 lr0 线性衰减到 lr0*lrf。--label-smoothing只在分类任务里生效检测任务里没有这个参数。注意不要同时把 lr0 调大和 label_smoothing 调大否则模型可能欠拟合。一次只改一个参数跑完对比验证集准确率再决定下一步。4. 模型评估与推理怎么验证 5 种花分类模型真的可用4.1 用验证集跑评估看懂混淆矩阵训练结束后runs/train-cls/flower_exp1目录下会有confusion_matrix.png和results.csv。混淆矩阵是最直观的工具能看出哪两类花容易混。比如 daisy 和 dandelion 都是黄色系小花模型可能把 dandelion 误判成 daisy。# 单独跑验证集评估 python classify/val.py \ --data flower_dataset \ --weights runs/train-cls/flower_exp1/weights/best.pt \ --batch-size 32 \ --imgsz 224 \ --device 0逻辑说明--weights指向训练保存的最佳权重通常是best.pt。评估会输出 top-1 和 top-5 准确率以及每个类别的精确率和召回率。如果某个类别召回率明显低说明这类样本被漏判得多需要补充这类样本或者检查标注。参数说明--batch-size评估时可以比训练时大因为不需要反向传播显存占用少。--imgsz必须和训练时一致否则准确率会掉。4.2 单张图片推理与批量推理验证完模型下一步就是实际推理。YOLOv5 分类推理命令很简洁# 单张图片推理 python classify/predict.py \ --weights runs/train-cls/flower_exp1/weights/best.pt \ --source test_images/rose_001.jpg \ --imgsz 224 # 批量推理整个目录 python classify/predict.py \ --weights runs/train-cls/flower_exp1/weights/best.pt \ --source test_images/ \ --imgsz 224 \ --save-txt逻辑说明--source可以是单张图片、目录或者 URL。--save-txt会把预测结果保存成 txt每行是类别和置信度。推理结果默认保存在runs/predict-cls/下。参数说明--imgsz同样要和训练一致。如果推理速度慢可以加--half用 FP16 推理但需要 GPU 支持。CPU 推理去掉--device参数即可。import torch from PIL import Image from torchvision import transforms # 加载模型 model torch.hub.load(ultralytics/yolov5, custom, pathruns/train-cls/flower_exp1/weights/best.pt) model.eval() # 预处理 transform transforms.Compose([ transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(), ]) img Image.open(test_images/rose_001.jpg) input_tensor transform(img).unsqueeze(0) # 推理 with torch.no_grad(): output model(input_tensor) prob torch.softmax(output, dim1) conf, pred torch.max(prob, dim1) print(f预测类别: {pred.item()}, 置信度: {conf.item():.4f})逻辑说明这段代码演示了如何把 YOLOv5 分类模型嵌入到 Python 脚本里。torch.hub.load加载自定义权重预处理要和训练时一致。softmax把输出转成概率torch.max取最大概率对应的类别。参数说明transforms.Resize(224)和训练时的 imgsz 保持一致。如果训练时用了数据增强推理时不需要增强只做 resize 和归一化。5. 避坑记录5 种花分类实战里最容易翻车的 4 个地方5.1 现象训练 loss 不下降acc1 一直卡在 0.2原因最常见的是数据目录结构不对。YOLOv5 分类要求 train 和 val 下直接是类别目录如果你多套了一层flower_dataset/train/images/daisy它就读不到标签。另一个原因是图片格式不支持比如 webp 或者灰度图虽然能读但可能出问题。解决用python classify/train.py --data flower_dataset之前先手动ls flower_dataset/train确认每个子目录都是类别名。图片统一转成 RGB jpg灰度图用 PIL 转一下。5.2 现象验证集准确率比训练集低 20 个点以上原因过拟合。5 种花数据集如果每类只有几百张模型很容易记住训练集。另外如果 train 和 val 的图片来自同一批拍摄分布太接近验证集就失去意义了。解决加--label-smoothing 0.1和--weight-decay 0.001同时把--epochs降到 50 左右。如果还不行检查 val 里的图片是不是和 train 有重复有的话重新切分。5.3 现象推理时预测结果全是同一类原因预处理不一致。训练时 YOLOv5 会做归一化均值和方差是 ImageNet 的如果你自己写推理脚本时忘了归一化或者 resize 方式不对模型输入分布就变了。解决用官方classify/predict.py先验证如果官方脚本正常再对比你的预处理代码。重点检查transforms.ToTensor()后面有没有加transforms.Normalize。5.4 现象GPU 显存够但训练速度很慢原因--batch-size太小或者--workers默认值太低。YOLOv5 分类默认 workers 是 8但在某些系统上会卡住反而拖慢速度。另外如果图片尺寸远大于 224dataloader 的 resize 会消耗大量 CPU。解决先把--batch-size提到显存允许的最大值再试--workers 4和--workers 0对比速度。如果图片原始尺寸很大提前用脚本批量 resize 到 256 左右减少 dataloader 压力。提示训练卡住不动时先看 GPU 利用率。如果利用率接近 0问题在数据加载如果利用率 100% 但 loss 不降问题在超参数或数据标签。6. 进阶技巧把 5 种花分类模型导出 ONNX 并验证一致性训练完模型下一步往往是部署。YOLOv5 分类模型可以导出 ONNX方便在树莓派5上部署自己训练的yolov5模型或者其他边缘设备上跑。导出命令python export.py \ --weights runs/train-cls/flower_exp1/weights/best.pt \ --include onnx \ --imgsz 224 \ --batch-size 1逻辑说明--include onnx指定导出格式--imgsz和--batch-size决定 ONNX 模型的输入维度。导出后的 onnx 文件默认和权重在同一目录。参数说明--batch-size 1适合边缘设备如果服务端批量推理可以设成 8 或 16。--imgsz必须和训练一致否则精度会掉。导出后一定要验证 ONNX 和 PyTorch 输出一致性import onnxruntime as ort import numpy as np import torch # 加载 ONNX 模型 sess ort.InferenceSession(runs/train-cls/flower_exp1/weights/best.onnx) input_name sess.get_inputs()[0].name # 构造随机输入 dummy np.random.randn(1, 3, 224, 224).astype(np.float32) # ONNX 推理 onnx_out sess.run(None, {input_name: dummy})[0] # PyTorch 推理 model torch.load(runs/train-cls/flower_exp1/weights/best.pt, map_locationcpu)[model] model.eval() with torch.no_grad(): torch_out model(torch.from_numpy(dummy)).numpy() # 对比最大差异 diff np.max(np.abs(onnx_out - torch_out)) print(f最大差异: {diff:.6f})逻辑说明ONNX Runtime 和 PyTorch 的数值计算有微小差异一般最大差异在 1e-4 以内算正常。如果超过 1e-3检查导出时的 opset 版本和输入尺寸。参数说明opset默认是 12如果部署环境支持更高版本可以调到 13 或 14。dummy输入的 shape 必须和导出时一致。我自己的习惯是每次导出 ONNX 后都跑一遍这个对比脚本差异超过 1e-3 就重新导出不将就。这个习惯帮我省了很多部署时的玄学问题。希望帮到你。本文还有配套的精品资源点击获取
返回列表