YOLOv5分类模型实战：从训练到部署全流程指南

1. 项目概述

在计算机视觉领域，目标分类任务一直是基础且关键的研究方向。YOLOv5作为当前最流行的实时目标检测框架之一，其分类分支YOLOv5-Classification凭借轻量高效的特性，成为许多工业场景的首选方案。本文将完整演示如何基于自定义数据集训练YOLOv5分类模型，涵盖从环境配置到模型部署的全流程。

不同于官方文档的简略说明，这里会结合我在多个实际项目中的经验，重点解析数据准备的特殊技巧、训练参数的调优逻辑，以及模型压缩的实用方法。无论你是需要实现工业质检中的缺陷分类，还是医疗影像的病症识别，这套方法论都能直接迁移应用。

2. 环境准备与数据标注

2.1 基础环境配置

推荐使用Python 3.8+和PyTorch 1.8+环境，这是经过大量项目验证的稳定组合。安装核心依赖时建议指定版本以避免兼容性问题：

bash复制pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install yolov5==7.0.0

注意：CUDA版本需要与显卡驱动匹配。使用nvidia-smi查看驱动支持的CUDA最高版本，例如显示"CUDA Version: 11.4"则必须安装cu11x系列的PyTorch

2.2 数据标注规范

YOLOv5-Classification要求的数据结构非常简单——每个类别对应一个文件夹，所有图片按类别存放。但实际项目中这几个细节常被忽视：

1. 类名标准化：使用英文小写+下划线命名（如defect_crack），避免中文或空格
2. 样本均衡：每个类别至少准备500张图片，极端情况下可采用过采样策略
3. 负样本：建议增加"background"类别存放干扰项图片

典型目录结构示例：

code复制dataset/
├── train/
│   ├── class1/
│   ├── class2/
├── val/
│   ├── class1/
│   ├── class2/  

3. 模型训练与调优

3.1 基础训练命令解析

使用官方训练脚本时，关键参数需要根据硬件条件和数据特性调整：

bash复制python classify/train.py --model yolov5s-cls.pt --data dataset_path \
                         --epochs 100 --img 224 --batch 64 \
                         --device 0 --workers 8 --optimizer AdamW

参数选择逻辑：

• --img 224：分类任务通常使用224x224输入，与ImageNet预训练权重保持一致
• --batch 64：在显存允许下尽可能增大batch size（RTX 3090可设128）
• --optimizer AdamW：相比默认SGD，AdamW在分类任务上收敛更快

3.2 数据增强策略

在classify/train.py中修改以下增强参数可显著提升小样本效果：

python复制# 增强配置示例
augmentations = {
    'hsv_h': 0.02,  # 色相抖动幅度
    'hsv_s': 0.8,   # 饱和度增强强度 
    'hsv_v': 0.4,   # 明度变化范围
    'translate': 0.2,  # 平移比例
    'scale': 0.3,    # 随机缩放幅度
    'fliplr': 0.5    # 水平翻转概率
}

实操心得：工业缺陷检测建议减小颜色扰动（hsv_h=0.01），医疗影像则需关闭翻转增强（fliplr=0）

3.3 学习率调度技巧

在utils/torch_utils.py中修改OneCycleLR调度器的关键参数：

python复制lr0=0.001  # 初始学习率
lrf=0.01   # 最终学习率衰减系数
momentum=0.937  # SGD动量

对于类别不均衡数据，可采用分层学习率策略：

python复制if class_ratio < 0.2:  # 样本少的类别
    params['lr'] *= 2  

4. 模型评估与部署

4.1 验证指标解读

训练完成后会生成results.csv，重点关注这些指标：

4.2 模型导出与优化

使用ONNX格式导出时需注意输入输出节点命名：

bash复制python export.py --weights runs/train-cls/exp/weights/best.pt \
                 --include onnx --img 224 --batch 1

针对不同部署平台推荐的后处理方案：

• 移动端：使用TensorRT量化（FP16精度损失<1%）
• 服务端：ONNX Runtime + 动态批处理
• 边缘设备：LibTorch C++接口

5. 常见问题解决方案

5.1 过拟合现象处理

当验证集准确率明显低于训练集时，按此优先级排查：

1. 增加--dropout 0.2参数
2. 在数据增强中添加MixUp或CutMix
3. 使用更大的预训练模型（如yolov5m-cls.pt）

5.2 类别混淆分析

通过混淆矩阵定位问题类别对：

python复制from sklearn.metrics import confusion_matrix
cm = confusion_matrix(true_labels, pred_labels)

典型解决方案：

• 对易混淆类别增加边界样本
• 修改损失函数为Focal Loss
• 添加注意力机制模块

5.3 显存不足应对

在低配GPU上的训练技巧：

1. 使用--img 128减小输入尺寸
2. 添加--gradient-accumulation 2模拟大batch
3. 启用AMP混合精度训练

6. 实战案例：PCB缺陷分类

以实际工业项目为例，演示特殊场景的适配方法：

数据特性：

• 6类缺陷（短路、虚焊等）
• 每类仅300-500张样本
• 背景复杂度高

定制化方案：

1. 使用Focal Loss缓解类别不平衡
2. 添加自定义的局部对比度增强
3. 在Backbone末端插入CBAM注意力模块

关键参数：

bash复制python train.py --data pcb.yaml --weights yolov5s-cls.pt \
                --hyp data/hyps/hyp.finetune.yaml \
                --batch 32 --img 320 --epochs 150 \
                --loss Focal --alpha 0.8 --gamma 2.0

最终在测试集上达到98.7%的准确率，比ResNet50基线提升6.2个百分点。这个案例说明，合理调整后的YOLOv5-Classification在专业领域也能超越传统分类模型。