★ FEATURED ARTICLE 简介本资源面向深度学习入门与计算机视觉实战人群提供一套可直接运行的YOLOv5图像分类方案用于解决五种花卉的识别分类问题。包内共2000个文件以1866张jpg花卉图片构成数据集主体并配有51个yaml配置文件、51个py脚本、12个yml、9个sh及若干md、ipynb、json等覆盖数据配置、模型定义、训练脚本与运行环境说明压缩包整体约257.86MB。资源已包含训练好的权重参数仅训练10个epochtop1准确度即达到0.91以上网络尚未完全收敛继续加大轮次可进一步提升性能。使用者只需按规范摆放datasets数据即可开始训练省去环境搭建与数据整理的重复工作。目前已有366人学习下载适合希望快速上手YOLOv5分类任务、验证模型效果或在此基础上做改进实验的读者参考。1. 从检测到分类YOLOV5 分类实战为什么值得单独拎出来讲很多人第一次接触 YOLOV5都是从画框、标 anchor、调 NMS 那套检测流程入门的于是下意识觉得 YOLOV5 只能做目标检测。但真到项目里你会发现有一类需求根本不需要框给你一张图只问「这是什么」比如 5 种花的分类数据集——雏菊、蒲公英、玫瑰、向日葵、郁金香输入一张图输出一个类别标签就够了。这时候再上检测头就是杀鸡用牛刀标注成本高、推理还慢。YOLOV5 官方仓库里其实一直带着一条分类分支classify/它把骨干换成更轻的结构、去掉检测头、接一个分类头训练和推理链路都是现成的。这篇笔记就围绕「YOLOV5 分类实战5 种花分类数据集」这条线把数据怎么摆、模型怎么选、超参数怎么调、训练完怎么验证、部署到边缘设备要注意什么一步步讲清楚。适合已经跑通过检测、想低成本切进分类任务的工程师也适合手上正好有个小数据集、想快速验证方案的人。2. YOLOV5 分类分支的目录结构与数据摆放规则2.1 为什么分类任务不直接复用检测权重先说选型理由。YOLOV5 的检测模型输出是(batch, anchors, 5classes)这种带框回归的张量分类分支输出是(batch, num_classes)的 logits两者头结构完全不同。你没法把yolov5s.pt的检测权重直接拿来当分类预训练权重用因为最后一层维度对不上强行加载只会报 shape mismatch。常见做法是分类任务用 ImageNet 预训练的骨干backbone权重或者干脆从头训。YOLOV5 分类分支支持--pretrained加载官方在 ImageNet 上训好的分类权重这对小数据集非常关键——5 种花加起来通常也就几千张图从头训很容易过拟合有预训练骨干收敛快得多。另一个容易踩的点是输入尺寸。检测默认640分类分支默认224这不是随便定的分类任务对全局语义敏感224 是 ImageNet 的标准输入骨干的下采样倍率也是按这个尺寸设计的。你要是硬改成 640特征图会大很多显存和耗时都上去了精度未必涨。2.2 数据集目录必须按 ImageFolder 规范摆YOLOV5 分类分支读数据用的是torchvision.datasets.ImageFolder这意味着目录结构必须严格是「一个类别一个文件夹」不能像检测那样用 txt 标注。5 种花的标准摆法如下# 数据集根目录结构train/val 各自独立 flower_dataset/ ├── train/ │ ├── daisy/ # 雏菊 │ │ ├── 001.jpg │ │ └── ... │ ├── dandelion/ # 蒲公英 │ ├── rose/ # 玫瑰 │ ├── sunflower/ # 向日葵 │ └── tulip/ # 郁金香 └── val/ ├── daisy/ ├── dandelion/ ├── rose/ ├── sunflower/ └── tulip/逻辑说明ImageFolder会扫描train/下的子目录名作为类别名按字典序排序生成class_to_idx。所以类别文件夹命名不要用中文也不要有空格否则后续推理时标签映射容易乱。参数上train和val必须物理隔离不能靠随机切分——分类任务里同一张图的不同增强版本如果同时出现在训练和验证集验证精度会虚高这是血泪经验。提示如果原始数据是一堆平铺的图片加一个 csv 标签文件先写个脚本按标签移动到对应文件夹别在训练脚本里临时切分容易出错且不可复现。2.3 用命令行跑通第一次训练数据摆好后进到 YOLOV5 仓库根目录分类训练入口是classify/train.py。最小可跑命令# 单卡训练5 种花分类输入 224 python classify/train.py \ --model yolov5s-cls.pt \ # 分类预训练权重不是检测的 yolov5s.pt --data ../flower_dataset \ # 指向含 train/val 的根目录 --epochs 50 \ --batch-size 32 \ --imgsz 224 \ --device 0逻辑说明--model指定骨干结构yolov5s-cls.pt是官方分类权重加载后只替换最后的全连接层为 5 类。--data指向的目录下必须有train和val两个子目录。--epochs 50对小数据集通常够用配合预训练权重一般 20 轮内就能看到验证精度趋稳。参数说明--batch-size在 8G 显存上 224 尺寸可以给到 64显存不够就往下调别硬撑导致 OOM。--imgsz保持 224除非你的花图分辨率极高且细节是分类关键。--device多卡写0,1单卡写0CPU 训练会慢到怀疑人生不推荐。3. 训练超参数怎么调从学习率到数据增强的实操参数3.1 学习率与优化器的搭配逻辑YOLOV5 分类分支默认用 SGD初始学习率lr00.01配合余弦退火。这个值在 ImageNet 这种百万级数据上没问题但 5 种花这种小数据集0.01偏大前几轮 loss 容易震荡甚至发散。我一般会把lr0降到0.001同时把--cos-lr打开让学习率平滑衰减。如果你用的是--pretrained加载的 ImageNet 权重更稳的做法是分两组参数骨干用小学习率比如lr0/10分类头用正常学习率。YOLOV5 分类脚本没有直接暴露分组学习率的参数常见做法是改classify/train.py里的 optimizer 构建部分或者干脆先用lr00.001整体微调等 loss 平稳后再看要不要解冻。优化器选择上SGD 泛化好但收敛慢Adam 收敛快但小数据集容易过拟合。我的习惯是数据量小于 5000 张用 Adamlr00.0005大于 5000 张用 SGDlr00.001。这不是铁律但能帮你少走几轮弯路。3.2 数据增强参数哪些该开哪些是坑YOLOV5 分类分支的增强参数集中在classify/train.py的--augment相关项常用的有--flip、--scale、--color、--erase。5 种花这个任务里花的朝向和颜色都是分类线索所以增强要克制参数默认值建议值理由flip0.50.5水平翻转对花分类无害可保留scale0.50.3缩放太狠会丢花瓣细节调小color0.50.2颜色是花的强特征别乱改色调erase0.00.1轻度随机擦除能抗遮挡别超过 0.2逻辑说明color增强会随机调整亮度、对比度、饱和度对花分类这种依赖颜色的任务开太大等于人为制造标签噪声。erase是随机遮挡一块区域适度开能提升模型对局部遮挡的鲁棒性但开太大可能把关键花瓣遮没反而掉点。注意增强参数不是越多越好。小数据集上增强过猛会让模型学不到真实分布验证精度上不去还找不到原因这是典型的玄学翻车现场。3.3 训练过程怎么盯看哪些指标、什么时候停训练日志里重点看三个数train_loss、val_loss、metrics/accuracy_top1。正常情况是 train_loss 持续下降val_loss 先降后平top1 精度稳步上升。如果 val_loss 开始上升而 train_loss 还在降就是过拟合信号该早停或加正则了。YOLOV5 分类脚本支持--patience参数做早停默认不启用。我一般设--patience 10意思是验证精度连续 10 轮不提升就停。这个值别设太小小数据集上精度波动大设 3、5 容易误停。另外--save-period控制 checkpoint 保存间隔默认只存最好的和最后一轮。如果你想分析训练过程可以设--save-period 5每 5 轮存一次方便回看哪一轮开始过拟合。4. 推理与验证把训练好的模型跑成可用的分类器4.1 单张图和批量推理的命令训练完权重默认在runs/train-cls/exp/weights/best.pt。推理入口是classify/predict.py# 单张图推理 python classify/predict.py \ --weights runs/train-cls/exp/weights/best.pt \ --source ../flower_dataset/val/rose/001.jpg \ --imgsz 224 # 整个文件夹批量推理结果存到 runs/predict-cls python classify/predict.py \ --weights runs/train-cls/exp/weights/best.pt \ --source ../flower_dataset/val \ --imgsz 224逻辑说明--source可以是单张图、文件夹、甚至视频流。批量推理时脚本会遍历文件夹下所有图片输出每张图的 top-5 类别和置信度。--imgsz必须和训练时一致否则精度会掉——这是很多人忽略的点训练 224 推理 320模型看到的特征尺度变了结果自然不准。参数说明--conf-thres控制置信度阈值分类任务里这个参数影响不大因为 softmax 输出总和为 1但可以设0.25过滤低置信样本。--save-txt会把结果写成 txt方便后续统计混淆矩阵。4.2 用验证集算混淆矩阵别只看 top1top1 精度只告诉你「对了多少」不告诉你「错在哪」。5 种花里玫瑰和郁金香在某些角度下容易混雏菊和蒲公英也有相似的花型。要定位问题得算混淆矩阵。YOLOV5 分类脚本本身不直接输出混淆矩阵常见做法是写个小脚本用val集跑推理把预测标签和真实标签对齐后交给 sklearnimport os import torch from torchvision import datasets, transforms from sklearn.metrics import confusion_matrix, classification_report # 加载模型和验证集 model torch.load(runs/train-cls/exp/weights/best.pt, map_locationcpu)[model] model.eval() val_dir ../flower_dataset/val transform transforms.Compose([ transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(), ]) val_set datasets.ImageFolder(val_dir, transformtransform) loader torch.utils.data.DataLoader(val_set, batch_size32, shuffleFalse) y_true, y_pred [], [] with torch.no_grad(): for imgs, labels in loader: outputs model(imgs) preds outputs.argmax(dim1) y_true.extend(labels.tolist()) y_pred.extend(preds.tolist()) # 输出混淆矩阵和每类指标 print(confusion_matrix(y_true, y_pred)) print(classification_report(y_true, y_pred, target_namesval_set.classes))逻辑说明这段脚本把验证集按 ImageFolder 读进来逐 batch 推理收集真实标签和预测标签最后用 sklearn 输出混淆矩阵和每类的 precision/recall/f1。target_names用val_set.classes自动对齐类别名避免手写错位。参数说明Resize(224)和CenterCrop(224)要和训练时的预处理一致训练用了 RandomResizedCrop验证就用 CenterCrop这是标准做法。batch_size按显存调推理阶段可以比训练大。拿到混淆矩阵后如果发现某两类互相错得特别多优先回去看这两类的训练样本是不是太少或太像而不是急着调模型结构。4.3 导出 ONNX 做跨平台部署训练验证都过了下一步往往是部署。YOLOV5 分类模型可以导出 ONNXpython export.py \ --weights runs/train-cls/exp/weights/best.pt \ --include onnx \ --imgsz 224 \ --batch-size 1逻辑说明--include onnx指定导出格式--batch-size 1是推理时的批大小部署到边缘设备通常用 1。导出后的 onnx 模型可以用 onnxruntime 加载也可以转成 TensorRT、RKNN 等格式。参数说明--imgsz必须和训练一致。如果目标设备只支持固定输入导出时就把 batch 和尺寸都固定死别用动态轴否则转换工具容易报错。导出后建议用 onnxruntime 跑一遍和 PyTorch 结果对比确认数值误差在可接受范围通常 1e-3 以内。5. 避坑与排查5 种花分类实战里最容易翻车的 5 个点5.1 现象训练 loss 不降精度卡在 20% 左右原因最常见的是--data路径指错或者train/val下没有按类别建文件夹ImageFolder 把整个目录当成一个类模型学不到区分性特征。另一个可能是--model误用了检测权重yolov5s.pt加载时虽然不报错但骨干初始化异常。解决先打印len(dataset.classes)确认类别数是 5不是 1。再确认--model用的是yolov5s-cls.pt。如果路径没问题把lr0降到0.0005再跑几轮看 loss 是否开始动。5.2 现象验证精度很高但实际推理一塌糊涂原因训练和推理的预处理不一致。训练时用了 RandomResizedCrop、ColorJitter 等增强推理时如果没做对应的归一化和尺寸对齐模型看到的输入分布和训练时差很远。解决推理脚本里的--imgsz必须等于训练时的--imgsz归一化参数mean/std也要一致。YOLOV5 分类默认用 ImageNet 的 mean/std如果你训练时改过推理也要同步改。5.3 现象某两类总是互相错分原因这两类的训练样本在视觉上太接近或者样本量严重不均衡。5 种花里玫瑰和郁金香在俯拍角度下确实容易混如果玫瑰样本只有 200 张而郁金香有 800 张模型会偏向郁金香。解决先看混淆矩阵确认是哪两类。如果是样本不均衡对少样本类做过采样或加 class weight。如果是视觉相似增加这两类的区分性样本或者用更强的骨干比如yolov5m-cls.pt提升特征表达能力。5.4 现象导出 ONNX 后推理结果和 PyTorch 对不上原因导出时--imgsz和训练不一致或者模型里有动态操作比如自适应池化在某些 opset 下行为不同。另外PyTorch 推理时如果忘了model.eval()BN 层会用 batch 统计量结果和 ONNX 的推理模式不一致。解决导出前先确认--imgsz和训练一致。导出后用 onnxruntime 跑同一张图和 PyTorch 输出逐元素对比误差超过 1e-2 就要查。PyTorch 侧记得model.eval()和torch.no_grad()。5.5 现象部署到树莓派或 RK3568 后帧率极低原因分类模型虽然比检测轻但yolov5s-cls在 ARM 上直接跑 PyTorch 仍然慢。没做量化、没转专用推理框架纯 CPU 推理一张 224 图可能要几百毫秒。解决先导出 ONNX再用目标平台的工具链转换和量化。RK3568 用 RKNN 工具链树莓派可以用 onnxruntime 或 NCNN。量化到 INT8 通常能提速 2-3 倍精度掉 1-2 个点以内可以接受。注意量化校准集要从训练集里抽别用验证集否则量化参数会过拟合。6. 进阶技巧用 TTA 和模型集成把 5 种花分类精度再抬一档训练跑通、部署上线之后如果精度还差那么一两个点别急着换模型先试两个成本极低的技巧TTA测试时增强和模型集成。TTA 的思路是推理时对同一张图做多种变换原图、水平翻转、多尺度缩放分别推理后把 softmax 概率平均。对花分类这种对翻转和轻微缩放不敏感的任务TTA 通常能稳定涨 1-2 个点。实现上不用改模型只在推理脚本里包一层import torch import torch.nn.functional as F def tta_predict(model, img_tensor): # img_tensor: (1, 3, 224, 224) probs [] # 原图 probs.append(F.softmax(model(img_tensor), dim1)) # 水平翻转 probs.append(F.softmax(model(torch.flip(img_tensor, dims[3])), dim1)) # 多尺度缩放到 256 再中心裁剪 224 img_256 F.interpolate(img_tensor, size256, modebilinear, align_cornersFalse) img_crop img_256[:, :, 16:240, 16:240] probs.append(F.softmax(model(img_crop), dim1)) # 平均概率 return torch.stack(probs).mean(dim0)逻辑说明这段函数对同一张图做三种视图推理把 softmax 后的概率平均。注意是平均概率不是平均 logits因为不同视图的 logits 尺度可能不同平均概率更稳。torch.flip的dims[3]是水平翻转对应宽度维度。参数说明多尺度的 256 和裁剪 224 是经验值你也可以试 288 裁 224。TTA 的代价是推理耗时变成 3 倍如果部署对延迟敏感可以只保留原图翻转两种涨点少一点但速度快一倍。模型集成更直接用不同骨干yolov5s-cls、yolov5m-cls或不同随机种子各训一个模型推理时平均概率。代价是训练和存储成本翻倍但精度通常比单模型高 2-3 个点。我的习惯是如果单模型已经 95% 以上TTA 就够了如果卡在 90% 上不去再考虑集成。最后说个我自己的教训小数据集上花在数据清洗和增强策略上的时间回报远高于调模型结构。我曾经在一个类似任务上换了三种骨干精度只涨了 0.5 个点后来发现训练集里有几十张标错的图清掉之后直接涨了 3 个点。所以每次精度上不去先回去看数据别急着改代码。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
