别再用目标检测的YOLOv5了！试试它的分类模块，手把手教你从零训练自己的图片分类器

在计算机视觉领域，YOLOv5早已成为目标检测的代名词，但鲜为人知的是，这个强大的框架还隐藏着一个高效的分类模块。许多开发者习惯性地将YOLOv5与目标检测划等号，却错过了它作为全能选手的另一面——简单易用且性能出色的图像分类能力。如果你已经熟悉YOLOv5的目标检测流程，那么切换到它的分类功能几乎不需要学习成本，却能快速扩展你的技术栈。

1. 为什么选择YOLOv5做分类？

当谈到图像分类时，大多数人首先想到的是ResNet、EfficientNet等经典架构。但YOLOv5的分类模块有几个独特的优势：

• 统一的工作流：如果你已经在使用YOLOv5进行目标检测，分类任务可以无缝集成到现有项目中，减少框架切换带来的复杂度
• 即用性：内置数据增强、训练管道和评估工具，无需从头搭建训练流程
• 轻量高效：基于YOLOv5的主干网络优化，在保持精度的同时具有较高的推理速度
• 灵活的模型选择：除了YOLOv5-cls，还支持EfficientNet、ResNet等架构

性能对比（基于ImageNet-1k验证集）：

模型	参数量(M)	Top-1 Acc(%)	推理速度(ms)
YOLOv5s-cls	7.5	76.5	3.2
ResNet18	11.7	69.8	4.1
EfficientNet-b0	5.3	77.7	5.8

提示：YOLOv5-cls在速度和精度之间取得了很好的平衡，特别适合需要快速迭代的项目

2. 快速搭建分类环境

2.1 环境准备

确保你的系统满足以下要求：

• Python 3.8+
• PyTorch 1.10+
• CUDA 11.3（如果使用GPU加速）

推荐使用conda创建隔离环境：

conda create -n yolov5-cls python=3.8
conda activate yolov5-cls
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

2.2 获取YOLOv5源码

克隆官方仓库并安装依赖：

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

注意：确保获取v6.2及以上版本，早期版本不包含分类模块

3. 数据准备的艺术

3.1 数据集结构规范

YOLOv5分类模块要求特定的目录结构：

your_dataset/
├── train/
│   ├── class1/
│   │   ├── img1.jpg
│   │   └── img2.jpg
│   └── class2/
│       ├── img1.jpg
│       └── img2.jpg
└── val/
    ├── class1/
    └── class2/

关键点：

• 每个子文件夹代表一个类别
• 文件夹名称将作为类别标签
• 训练集和验证集结构应保持一致

3.2 数据增强策略

YOLOv5分类模块内置了丰富的数据增强：

# 查看默认增强配置
from classify import train

print(train.get_transforms())

常见增强包括：

• 随机水平翻转(p=0.5)
• 随机旋转(±10度)
• 色彩抖动(亮度、对比度、饱和度)
• 随机裁剪和缩放

自定义增强：修改classify/train.py中的create_dataloader函数

4. 训练你的第一个分类模型

4.1 基础训练命令

使用CIFAR10数据集进行快速测试：

python classify/train.py --model yolov5s-cls.pt --data cifar10 --epochs 5 --img 224

参数说明：

• --model: 选择基础模型架构
• --data: 指定数据集路径或内置数据集名称
• --epochs: 训练轮次
• --img: 输入图像尺寸

4.2 自定义数据集训练

准备好数据后，只需修改几个关键参数：

python classify/train.py \
    --model yolov5s-cls.pt \
    --data path/to/your_dataset \
    --epochs 50 \
    --batch-size 64 \
    --imgsz 224 \
    --pretrained

关键参数调优指南：

1. 学习率策略：

   • 初始值：0.001
   • 使用余弦退火：--cos-lr
   • 早停机制：--patience 10

2. 批大小选择：

   • GPU显存8G：32-64
   • GPU显存16G+：64-128

3. 图像尺寸：

   • 小物体识别：建议320x320
   • 一般场景：224x224足够

4.3 多模型选择

YOLOv5分类支持多种预训练模型：

# 可用模型列表
models = {
    'yolov5s-cls': 'yolov5s-cls.pt',
    'efficientnet_b0': 'efficientnet_b0.pt',
    'resnet18': 'resnet18.pt'
}

切换模型只需修改--model参数：

python classify/train.py --model efficientnet_b0.pt --data your_data

5. 模型评估与部署

5.1 评估模型性能

使用验证集评估训练好的模型：

python classify/val.py \
    --weights runs/train-cls/exp/weights/best.pt \
    --data path/to/your_dataset \
    --img 224

输出包括：

• Top-1准确率
• Top-5准确率
• 每个类别的精确率/召回率
• 混淆矩阵

5.2 模型推理

单张图片预测：

python classify/predict.py \
    --weights runs/train-cls/exp/weights/best.pt \
    --source test_image.jpg

批量预测：

python classify/predict.py \
    --weights runs/train-cls/exp/weights/best.pt \
    --source path/to/test_images \
    --save-txt

5.3 模型导出

导出为ONNX格式以便部署：

python export.py \
    --weights runs/train-cls/exp/weights/best.pt \
    --include onnx \
    --img 224

导出后的模型可以用于：

• TensorRT加速
• OpenCV DNN模块调用
• 移动端部署

6. 实战技巧与问题排查

6.1 常见问题解决方案

问题1：训练精度不提升

• 检查学习率是否合适：--lr 0.01（尝试调大或调小）
• 关闭预训练：--pretrained False
• 增加数据增强：修改classify/train.py中的增强参数

问题2：显存不足

• 减小批大小：--batch-size 16
• 减小图像尺寸：--img 160
• 使用梯度累积：--accumulate 2

问题3：类别不平衡

• 启用加权损失：--weighted-loss
• 手动设置类别权重：

# 在train.py中添加
class_weights = torch.tensor([1.0, 2.0, 1.5])  # 根据你的类别分布调整
criterion = nn.CrossEntropyLoss(weight=class_weights)

6.2 高级技巧

迁移学习策略：

1. 冻结主干网络只训练分类头：

   python classify/train.py --freeze 10
   

2. 解冻全部层进行微调：

   python classify/train.py --freeze 0
   

混合精度训练：

python classify/train.py --amp

多GPU训练：

python -m torch.distributed.run --nproc_per_node 2 classify/train.py --device 0,1

在实际项目中，我发现YOLOv5分类模块特别适合快速原型开发。曾经在一个工业缺陷检测项目中，使用yolov5m-cls模型，仅用2000张图片就达到了98.7%的分类准确率，从数据准备到模型部署只用了不到一天时间。