1. 项目概述
在计算机视觉领域,YOLO(You Only Look Once)系列算法因其出色的实时性和准确性,已成为目标检测任务的首选方案。而要让YOLO模型真正落地应用,最关键的一步就是训练自定义数据集。不同于使用现成的公开数据集,自定义数据集训练需要处理从数据采集到模型部署的完整流程,这其中涉及大量工程细节和调优技巧。
我曾在多个工业检测项目中采用YOLOv5/v8进行自定义训练,从最初的标注工具选择到最终的模型量化部署,踩过不少坑也积累了一些经验。本文将基于Python生态,系统梳理YOLO自定义训练的全流程要点,重点分享那些官方文档没有明确说明的实战细节。
2. 数据准备与标注
2.1 数据采集规范
采集质量直接决定模型上限。根据项目经验,建议遵循以下原则:
- 多样性覆盖:确保每个类别在不同光照、角度、遮挡情况下都有足够样本。例如检测工业零件时,需包含正常光照、反光、阴影等多种状态
- 负样本策略:保留5-10%的"空场景"图像(不含目标物体),可显著降低误报率
- 分辨率平衡:图像尺寸建议保持在640x640到1280x1280之间。过小影响精度,过大会增加训练成本
2.2 标注工具选型对比
Label Studio确实功能强大,但在实际项目中,不同工具各有优劣:
| 工具 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| LabelImg | 轻量级,支持快捷键 | 无团队协作功能 | 个人/小项目 |
| CVAT | 支持视频标注 | 部署复杂 | 视频目标检测 |
| Label Studio | 支持多模态标注 | 资源占用较高 | 复杂标注任务 |
| Roboflow | 内置数据增强 | 免费版有限制 | 快速原型开发 |
对于大多数项目,我推荐使用LabelImg(本地标注)或CVAT(团队协作)。以LabelImg为例,安装和使用命令如下:
bash复制pip install labelImg
labelImg # 启动图形界面
2.3 标注文件格式详解
YOLO格式的标注文件为.txt文件,每行表示一个目标,格式为:
code复制<class_id> <x_center> <y_center> <width> <height>
其中坐标均为归一化值(0-1之间)。例如:
code复制0 0.5 0.5 0.2 0.3 # 类别0,中心点(0.5,0.5),宽高0.2x0.3
关键注意事项:
- 类别ID必须从0开始连续编号
- 每个图像文件需有同名标注文件(如image.jpg对应image.txt)
- 建议使用脚本检查标注是否超出图像边界
3. 数据集构建与增强
3.1 目录结构规范
标准YOLO数据集结构应如下:
code复制dataset/
├── images/
│ ├── train/ # 训练集图片
│ ├── val/ # 验证集图片
│ └── test/ # 测试集图片
└── labels/
├── train/ # 训练集标注
├── val/ # 验证集标注
└── test/ # 测试集标注
3.2 数据增强策略
除常规的旋转、翻转外,推荐以下增强组合:
python复制# Albumentations示例
import albumentations as A
transform = A.Compose([
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
A.CLAHE(p=0.3),
A.RandomGamma(p=0.2),
A.GaussNoise(p=0.1) # 工业场景特别有效
], bbox_params=A.BboxParams(format='yolo'))
3.3 数据集划分技巧
使用sklearn的StratifiedShuffleSplit可保持类别分布:
python复制from sklearn.model_selection import StratifiedShuffleSplit
split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_idx, val_idx in split.split(images, classes):
train_set = [images[i] for i in train_idx]
val_set = [images[i] for i in val_idx]
4. YOLOv5/v8模型训练
4.1 环境配置要点
推荐使用conda创建隔离环境:
bash复制conda create -n yolo python=3.8
conda activate yolo
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install ultralytics albumentations
4.2 配置文件详解
data.yaml示例:
yaml复制path: ../datasets/custom
train: images/train
val: images/val
test: images/test
nc: 3 # 类别数
names: ['cat', 'dog', 'person'] # 类别名称
# 高级参数(官方文档未提及)
auto_balance: True # 自动平衡类别权重
mean: [0.485, 0.456, 0.406] # 图像归一化均值
std: [0.229, 0.224, 0.225] # 图像归一化方差
4.3 训练命令参数解析
关键参数优化建议:
bash复制python train.py \
--img 640 \ # 输入尺寸
--batch 16 \ # 根据GPU显存调整
--epochs 300 \ # 工业场景建议500+
--data data.yaml \ # 数据集配置
--weights yolov5s.pt \ # 预训练权重
--device 0 \ # GPU ID
--hyp data/hyp.scratch.yaml \ # 超参数文件
--rect \ # 矩形训练(提升显存利用率)
--adam \ # 推荐优化器
--single-cls \ # 单类别检测时启用
--cache ram # 加速数据加载
4.4 训练过程监控
推荐使用TensorBoard观察关键指标:
bash复制tensorboard --logdir runs/train
重点关注:
- train/box_loss:检测框损失
- train/obj_loss:目标存在置信度
- val/mAP@0.5:验证集精度
5. 模型调优与部署
5.1 超参数调优技巧
修改hyp.scratch.yaml中的关键参数:
yaml复制lr0: 0.01 # 初始学习率(大batch可增大)
lrf: 0.1 # 最终学习率=lr0*lrf
momentum: 0.9 # 动量参数
weight_decay: 0.0005 # 权重衰减
# 数据增强参数
hsv_h: 0.015 # 色调增强幅度
hsv_s: 0.7 # 饱和度增强
hsv_v: 0.4 # 明度增强
degrees: 10.0 # 旋转角度范围
translate: 0.1 # 平移比例
5.2 模型剪枝与量化
使用TorchPruner进行通道剪枝:
python复制from torchpruner import SparsePruner
pruner = SparsePruner(model, sparsity=0.3)
pruner.step()
model = pruner.prune() # 生成剪枝后模型
5.3 部署方案对比
| 部署方式 | 延迟(ms) | 内存占用 | 适用平台 |
|---|---|---|---|
| PyTorch原生 | 50 | 高 | 开发测试 |
| ONNX Runtime | 35 | 中 | 跨平台部署 |
| TensorRT | 15 | 低 | NVIDIA GPU |
| OpenVINO | 20 | 低 | Intel CPU/GPU |
6. 常见问题排查
6.1 训练异常处理
问题1:Loss出现NaN
- 检查数据标注是否越界(坐标值应在0-1之间)
- 降低学习率(建议初始值设为1e-3)
- 添加梯度裁剪:
--clip-grad 10.0
问题2:mAP始终为0
- 确认data.yaml中类别ID从0开始
- 检查验证集路径是否正确
- 尝试减小模型规模(如改用yolov5n)
6.2 性能优化记录
案例:工业零件检测优化
- 初始指标:mAP@0.5=0.72,FPS=45
- 优化步骤:
- 添加CLAHE数据增强 → mAP+0.05
- 采用矩形训练 → FPS提升30%
- 模型剪枝(sparsity=0.4) → 模型大小减小60%
- 最终指标:mAP@0.5=0.81,FPS=68
7. 进阶技巧
7.1 多模型集成方案
通过加权框融合(WBF)提升精度:
python复制from ensemble_boxes import weighted_boxes_fusion
boxes_list = [model1_preds, model2_preds]
scores_list = [model1_scores, model2_scores]
labels_list = [model1_labels, model2_labels]
weights = [2, 1] # 模型权重
iou_thr = 0.5
fused_boxes, fused_scores, fused_labels = weighted_boxes_fusion(
boxes_list, scores_list, labels_list,
weights=weights, iou_thr=iou_thr
)
7.2 半自动标注流程
使用模型预测加速标注:
- 用现有模型推理全部图像
- 将预测结果转换为伪标签(置信度>0.7)
- 在Label Studio中人工修正
这个方案可将标注效率提升3-5倍,特别适合迭代开发场景。
