YOLOV5分类实战：花分类数据集训练与推理全流程

简介本资源面向深度学习入门与计算机视觉实战人群提供一套可直接运行的YOLOv5图像分类方案用于解决五种花卉的识别分类问题。包内共2000个文件以1866张jpg花卉图片构成数据集主体并配有51个yaml配置、51个py脚本、12个yml环境文件及9个sh运行脚本另有ipynb教程笔记与md说明文档压缩包整体约257.86MB目录结构清晰摆放好datasets数据即可开始训练。资源已附带训练好的权重参数仅训练10个epochtop1准确度即达到0.91以上网络尚未完全收敛加大轮次可进一步提升性能。目前已有366人学习下载适合希望快速跑通分类流程、理解YOLOv5分类任务配置与训练细节的读者参考也可作为课程设计或小型项目的实践素材。1. 花分类数据集实战从 257MB 资源包到 0.91 top1 的落地路径拿到一个 257MB 的 YOLOV5 分类资源包里面塞着代码、数据集和训练好的权重很多人第一反应是「直接跑起来看看」。但真正上手前得先想清楚一件事YOLOV5 做分类和做检测是两条完全不同的链路检测那套 anchor、NMS、后处理在这里统统用不上分类只关心一件事——这张图属于哪一类。这份资源包解决的就是「不想从零搭分类框架」的问题5 种花的图片已经按目录分好训练脚本和推理代码都在权重也给你了10 个 epoch 跑出 top1 0.91 的成绩。适合谁手头有分类需求但不想折腾数据划分和训练循环的从业者或者想拿一个能跑通的分类 baseline 快速验证自己想法的工程师。资源包里那几张 jpg 就是数据集样本tutorial.ipynb 是入口.gitattributes 说明这包是从 git 仓库导出的结构上没毛病。2. YOLOV5 分类链路拆解为什么选它而不是 ResNet2.1 分类任务里 YOLOV5 的真实定位YOLOV5 在检测圈出名但它的分类分支其实是一个独立的轻量级 backbone 加分类头。资源包里用的就是这套结构不是把检测头拆掉改的。为什么有人拿它做分类两个原因一是工程化程度高数据加载、增强、学习率调度、混合精度这些训练细节都封装好了你只需要把图片按类别放进文件夹二是模型体积小YOLOV5s 分类版本的参数量在几百万级别推理速度快适合后续往边缘设备上搬。但要注意分类精度上它不一定打得过专门为分类设计的 ResNet 或 EfficientNet尤其是数据量大的时候。这份资源包 10 个 epoch 到 0.91说明数据本身区分度不错花类之间的视觉差异比较明显换个难一点的数据集同样的 epoch 数可能连 0.7 都费劲。2.2 数据集目录结构与加载逻辑资源包里的图片文件名是一串数字加哈希这是 Flickr 那类图片站导出的典型命名。分类任务不看文件名看的是它所在的文件夹名。YOLOV5 分类的数据加载器会扫描 datasets 目录下的子文件夹每个子文件夹名就是一个类别标签。所以摆放数据的时候目录结构必须是这样的datasets/ ├── train/ │ ├── daisy/ │ │ ├── 4932735362_6e1017140f.jpg │ │ └── 4341530649_c17bbc5d01.jpg │ ├── rose/ │ │ ├── 5693459303_e61d9a9533.jpg │ │ └── 10164073235_f29931d91e.jpg │ ├── sunflower/ │ │ ├── 6140892289_92805cc590.jpg │ │ └── 244074259_47ce6d3ef9.jpg │ ├── tulip/ │ │ └── 3704306975_75b74497d8.jpg │ └── dandelion/ │ └── 6140693467_211a135b6d.jpg └── val/ ├── daisy/ ├── rose/ ├── sunflower/ ├── tulip/ └── dandelion/训练集和验证集要分开比例常见做法是 8:2 或 7:3。资源包里如果已经分好了直接确认一下 val 目录下每个类别都有图片别出现某个类验证集为空的情况否则训练日志里那个类的准确率会直接报 nan。图片格式支持 jpg、png、bmp 这些常见的尺寸不要求统一加载器会做 resize默认是 224x224这个参数在训练脚本里可以改。2.3 训练入口与关键参数资源包里的训练脚本一般叫classify/train.py调用方式是这样python classify/train.py \ --model yolov5s-cls.pt \ --data datasets/ \ --epochs 10 \ --batch-size 32 \ --imgsz 224 \ --device 0 \ --project runs/train-cls \ --name flower_exp逐项说清楚--model指定 backboneyolov5s-cls.pt是分类专用的预训练权重不是检测那个--data指向数据集根目录脚本会自动去找 train 和 val 子目录--epochs 10就是资源包作者用的轮次想涨点就往上加30 到 50 是常见区间--batch-size 32看显存调8G 显存跑 224 尺寸一般能到 64--imgsz 224是输入分辨率花分类这个尺寸够用想更细可以上 320但训练时间会涨--device 0指定第一块 GPU没 GPU 就写 cpu但速度会慢到你想砸键盘。训练过程中会打印每个 epoch 的 loss 和 top1、top5 准确率10 个 epoch 后看最后一行的 top1 是不是在 0.91 附近如果差太多先检查数据目录有没有摆错。2.4 推理验证拿权重跑单张图训练完权重会存在runs/train-cls/flower_exp/weights/best.pt推理脚本这样调python classify/predict.py \ --weights runs/train-cls/flower_exp/weights/best.pt \ --source datasets/val/daisy/4932735362_6e1017140f.jpg \ --imgsz 224输出会显示 top1 类别和置信度。如果资源包里已经带了训练好的权重直接拿它跑一遍验证集看看准确率是不是和作者说的一致。常见做法是写个小循环遍历 val 下所有图片统计预测正确的比例这样比单张看更靠谱。注意推理时的--imgsz要和训练时一致不一致的话精度会掉这是血泪经验别问我是怎么知道的。3. 从零复现训练数据摆放、参数调整与结果验证3.1 数据划分的坑与正确做法资源包里的数据可能已经分好了 train 和 val但如果你要换自己的数据集划分这一步最容易翻车。常见错误是随机打乱所有图片再按比例切这样会导致同一个物体的不同角度图片同时出现在训练集和验证集里验证准确率虚高。正确做法是按类别分层抽样每个类里单独切出 20% 做验证。用 Python 写就是import os import random import shutil src_dir datasets/all_images dst_train datasets/train dst_val datasets/val split_ratio 0.8 for cls in os.listdir(src_dir): cls_path os.path.join(src_dir, cls) imgs os.listdir(cls_path) random.shuffle(imgs) split_idx int(len(imgs) * split_ratio) train_imgs imgs[:split_idx] val_imgs imgs[split_idx:] os.makedirs(os.path.join(dst_train, cls), exist_okTrue) os.makedirs(os.path.join(dst_val, cls), exist_okTrue) for img in train_imgs: shutil.copy(os.path.join(cls_path, img), os.path.join(dst_train, cls, img)) for img in val_imgs: shutil.copy(os.path.join(cls_path, img), os.path.join(dst_val, cls, img))这段代码按类别分别划分保证每个类在验证集里都有代表。split_ratio控制训练集比例0.8 意味着 80% 训练、20% 验证。如果某个类图片特别少比如只有 10 张那验证集就 2 张统计意义不大这时候要么多找数据要么用交叉验证。3.2 超参数怎么调学习率、batch size 与 epochYOLOV5 分类的默认学习率是 0.001优化器是 SGD动量 0.9。这个配置在花分类这种小数据集上一般能跑但如果你发现 loss 震荡得厉害先把学习率降到 0.0005 试试。batch size 前面说了看显存但别设太小小于 16 的时候 BN 层的统计量会不准精度会掉。epoch 数资源包用了 10作者也说了网络还没收敛加大轮次能涨点。我一般会先跑 10 个 epoch 看趋势如果验证准确率还在涨就加到 30 或 50同时开早停patience 设 10连续 10 个 epoch 验证准确率不涨就停省时间。python classify/train.py \ --model yolov5s-cls.pt \ --data datasets/ \ --epochs 50 \ --batch-size 32 \ --imgsz 224 \ --lr0 0.001 \ --patience 10 \ --device 0--lr0是初始学习率--patience是早停耐心值。这两个参数配合好了既能涨点又不浪费时间。注意 YOLOV5 的学习率调度是带 warmup 的前几个 epoch 学习率会从很小慢慢升上去所以前 3 个 epoch 准确率低是正常的别急着停。3.3 训练过程监控与结果解读训练日志里重点看三个东西训练 loss、验证 loss、top1 准确率。训练 loss 一直降但验证 loss 开始涨说明过拟合了要么加数据增强要么减模型复杂度。top1 准确率在验证集上算的资源包说 0.91你跑出来如果差很多先确认验证集图片有没有问题。常见做法是训练结束后用混淆矩阵看哪些类容易混花分类里玫瑰和郁金香在某些角度下确实像如果这两类互相错得多可以考虑加更多这两类的训练数据。from sklearn.metrics import confusion_matrix import seaborn as sns import matplotlib.pyplot as plt # 假设你已经拿到了真实标签和预测标签 cm confusion_matrix(true_labels, pred_labels) sns.heatmap(cm, annotTrue, fmtd, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(Predicted) plt.ylabel(True) plt.show()这段代码画混淆矩阵true_labels和pred_labels是你推理完验证集后得到的列表class_names是类别名列表。看矩阵对角线越深越好非对角线上的数字就是错分的数量。如果某个类错得特别多针对性补数据比盲目加 epoch 有效。4. 避坑与排查分类训练里那些让人抓狂的瞬间4.1 验证集准确率不动或为 0现象训练跑起来了loss 也在降但验证集 top1 一直是 0 或者一个固定值。原因通常是验证集目录结构不对比如 val 下面直接放图片而不是按类别放子文件夹加载器找不到类别标签所有验证图片被当成一个类。解决检查 val 目录确保每个类别一个子文件夹且子文件夹里至少有图片。另一个可能是--data参数指错了指向了 train 的父目录而不是包含 train 和 val 的根目录。4.2 显存溢出CUDA out of memory现象训练刚开始就报 CUDA out of memory。原因batch size 太大或者 imgsz 太大。解决先把 batch size 减半如果还不行就降 imgsz从 224 降到 192 或 160。另外检查有没有其他进程占着显存nvidia-smi看一眼。YOLOV5 分类模型本身不大8G 显存跑 224 尺寸 batch 64 一般没问题如果报错先怀疑是不是同时跑了别的任务。4.3 训练 loss 变成 nan现象几个 epoch 后 loss 突然变成 nan。原因学习率太大导致梯度爆炸或者数据里有损坏的图片。解决先把学习率降到 0.0001 跑几个 epoch 看看如果正常了再慢慢往上调。数据方面用 PIL 遍历一遍所有图片能打开且能 convert 成 RGB 的才是好图损坏的直接删掉。from PIL import Image import os for root, dirs, files in os.walk(datasets): for f in files: if f.endswith((.jpg, .png, .bmp)): try: img Image.open(os.path.join(root, f)) img.convert(RGB) except: print(坏图:, os.path.join(root, f))这段代码扫一遍数据集把打不开的图片揪出来。坏图不删训练时加载到就会报错或者产生 nan。4.4 推理结果和训练准确率对不上现象训练时验证集准确率 0.91但自己拿几张图推理错得离谱。原因推理时的预处理和训练时不一致比如训练用了归一化而推理没做或者 imgsz 不一致。解决直接用 YOLOV5 的 predict 脚本它内部会做和训练一致的预处理。如果自己写推理代码记得把图片 resize 到训练时的尺寸然后除以 255 归一化再转成 tensor。4.5 类别不平衡导致某些类全错现象混淆矩阵里某个类的准确率特别低几乎全错到另一个类。原因训练数据里这个类的图片太少模型没学够。解决要么补数据要么在 loss 里加类别权重让少样本类的 loss 权重大一些。YOLOV5 分类默认不带类别权重需要自己改 loss 函数常见做法是用nn.CrossEntropyLoss(weightclass_weights)class_weights按类别频率的倒数算。5. 进阶技巧用这份资源包做迁移学习和模型导出资源包里的权重不只是拿来跑花分类的它本身就是一个在 ImageNet 上预训练过的 backbone你可以拿它做迁移学习换自己的数据集只训练分类头或者微调整个网络。具体做法是把--model指向资源包里的best.pt然后--data指向你自己的数据集--epochs设小一点10 到 20 就够因为 backbone 已经学过了你只需要它适应新类别。学习率也要调小用 0.0001 或者 0.0005避免把预训练权重冲掉。python classify/train.py \ --model runs/train-cls/flower_exp/weights/best.pt \ --data datasets/my_custom_data/ \ --epochs 15 \ --batch-size 32 \ --imgsz 224 \ --lr0 0.0005 \ --device 0这里--model指向的是花分类训练好的权重不是原始的 yolov5s-cls.pt。这样做的效果通常比从 ImageNet 权重开始好因为花分类任务学到的特征和很多细粒度分类任务更接近。但要注意如果你的新数据集和花差异很大比如工业缺陷检测那还是用原始 ImageNet 权重更稳。另一个进阶用法是导出模型给推理引擎用。YOLOV5 支持导出 ONNX、TorchScript 等格式导出命令是python export.py \ --weights runs/train-cls/flower_exp/weights/best.pt \ --include onnx \ --imgsz 224导出的 ONNX 模型可以丢给 ONNX Runtime 或者 TensorRT 跑推理速度比 PyTorch 原生快不少。注意导出时的--imgsz要和训练时一致否则精度会掉。导出后拿一张验证集图片跑一下 ONNX 推理和 PyTorch 的结果对比误差在 1e-3 以内算正常。验证导出模型是否正确我一般会写个小脚本对比 PyTorch 和 ONNX 的输出import torch import onnxruntime as ort import numpy as np from PIL import Image # PyTorch 推理 model torch.load(runs/train-cls/flower_exp/weights/best.pt, map_locationcpu)[model] model.eval() img Image.open(datasets/val/daisy/4932735362_6e1017140f.jpg).resize((224, 224)) img_tensor torch.tensor(np.array(img)).permute(2, 0, 1).float().unsqueeze(0) / 255.0 with torch.no_grad(): pt_out model(img_tensor).numpy() # ONNX 推理 sess ort.InferenceSession(runs/train-cls/flower_exp/weights/best.onnx) onnx_out sess.run(None, {images: img_tensor.numpy()})[0] print(PyTorch 输出:, pt_out.argmax(), pt_out.max()) print(ONNX 输出:, onnx_out.argmax(), onnx_out.max()) print(最大误差:, np.abs(pt_out - onnx_out).max())这段代码分别用 PyTorch 和 ONNX Runtime 跑同一张图对比类别和置信度。如果最大误差在 1e-3 以内说明导出没问题。如果误差大检查导出时的 opset 版本和输入名称ONNX 的输入名一般是images但不同版本可能不一样用sess.get_inputs()[0].name确认一下。从那以后我每次拿到一个训练好的分类权重都会先跑一遍验证集确认精度再导出 ONNX 对比输出最后才往业务里集成。这个习惯帮我省了很多「训练时好好的上线就翻车」的后悔药。希望帮到你。本文还有配套的精品资源点击获取