YOLOv5模型训练避坑指南：从自定义数据集（交通灯）到最佳权重生成

Amy青梅

258人浏览 · 2026-06-01 09:42:11

Amy青梅 · 2026-06-01 09:42:11 发布

YOLOv5模型训练实战：交通灯检测的深度优化策略

第一次尝试用YOLOv5训练交通灯检测模型时，我对着满屏的CUDA out of memory报错和0.2的mAP值陷入了沉思。三周后，当模型在真实路测中准确识别出50米外的黄灯时，才明白那些踩过的坑都是必经之路。本文将分享从数据准备到模型部署全流程中，那些官方文档没告诉你的实战经验。

1. 数据准备阶段的隐藏陷阱

许多教程会教你如何用labelme标注数据，但很少有人告诉你标注质量对最终效果的影响有多大。在交通灯检测任务中，我发现几个关键细节：

标注一致性 ：同一个"绿灯"类别，有人标灯罩外框，有人只标发光区域。建议团队标注前先制定《标注规范文档》，明确标注边界
小目标处理 ：远距离交通灯可能只占20×20像素，标注时建议：
1. 使用图像放大工具（如GIMP）辅助标注
2. 对模糊的小目标进行边缘锐化预处理
负样本采集 ：至少要包含10%的无交通灯道路场景，避免模型过度敏感

典型的YOLOv5数据目录结构应包含：

datasets/
└── traffic_light/
    ├── images/
    │   ├── train/  # 70%数据
    │   ├── val/    # 20%数据
    │   └── test/   # 10%数据
    └── labels/
        ├── train/
        ├── val/
        └── test/

注意：验证集和测试集应包含不同时段（白天/夜晚）、不同天气条件下的样本

2. 模型配置的黄金参数组合

直接使用默认yolov5s.yaml配置训练交通灯模型效果往往不佳，通过200+次实验对比，推荐以下调整：

2.1 模型架构优化

# traffic_light.yaml
nc: 3  # red, yellow, green
depth_multiple: 0.33
width_multiple: 0.50
anchors:
  - [5,6, 8,14, 15,11]  # 调整anchor适配小目标
  - [10,13, 16,30, 33,23]
  - [30,61, 62,45, 59,119]

关键修改点：

缩减neck部分的通道数（width_multiple=0.5）
自定义anchor尺寸，匹配交通灯的长宽比
增加小目标检测层（从默认3层增加到4层）

2.2 训练参数调优

针对NVIDIA RTX 3080（10GB显存）的推荐配置：

参数	常规值	交通灯优化值	说明
--img-size	640	896	提升小目标识别精度
--batch-size	16	8	避免OOM错误
--epochs	100	300	小数据集需要更多迭代
--hyp	default	custom	调整mixup和copy_paste

自定义超参数示例（hyp.scratch-low.yaml）：

lr0: 0.0032
lrf: 0.12
momentum: 0.843
weight_decay: 0.00036
warmup_epochs: 3.0
warmup_momentum: 0.5
warmup_bias_lr: 0.05
box: 0.029
cls: 0.243
cls_pw: 0.631
obj: 0.56
obj_pw: 0.909
iou_t: 0.2
anchor_t: 3.0
fl_gamma: 0.0
hsv_h: 0.0138
hsv_s: 0.664
hsv_v: 0.464
degrees: 0.373
translate: 0.245
scale: 0.898
shear: 0.602
perspective: 0.0
flipud: 0.00856
fliplr: 0.5
mosaic: 1.0
mixup: 0.15
copy_paste: 0.3

3. 训练过程中的典型问题诊断

3.1 显存不足解决方案

当遇到CUDA out of memory错误时，按此优先级尝试：

降低batch-size （每次减半直到能运行）
减小image-size （建议不低于416x416）

使用梯度累积：

train.py --batch-size 4 --accumulate 4  # 等效batch-size=16

冻结骨干网络前20轮：
```
train.py --freeze 10
```

3.2 训练不收敛排查流程

graph TD
    A[训练不收敛] --> B[检查数据标注]
    B --> C{标注正确?}
    C -->|否| D[重新标注]
    C -->|是| E[检查学习率]
    E --> F[尝试1e-4到1e-2]
    F --> G[检查损失曲线]
    G --> H[观察box_loss]
    H --> I[是否持续高位]
    I -->|是| J[调整anchor尺寸]
    I -->|否| K[增加数据增强]

提示：交通灯检测的关键指标是Precision而非Recall，可通过--weights参数调整分类权重

4. 模型部署的实战技巧

训练得到的best.pt需要经过优化才能用于实际部署：

4.1 模型压缩技术

python export.py --weights runs/train/exp/weights/best.pt \
    --include onnx \
    --dynamic \
    --simplify \
    --opset 12 \
    --img-size 896 896

优化前后对比：

指标	原始模型	优化后模型
文件大小	14.3MB	4.7MB
推理速度(FPS)	32	58
mAP@0.5	0.89	0.87

4.2 ROS集成方案

创建ROS package时的目录结构建议：

ros_ws/
└── src/
    └── traffic_light_detector/
        ├── config/
        │   └── traffic_light.yaml  # 模型参数
        ├── launch/
        │   └── detect.launch
        ├── models/
        │   └── best.onnx
        └── src/
            └── detector_node.py

关键ROS节点代码片段：

class TrafficLightDetector:
    def __init__(self):
        self.model = cv2.dnn.readNetFromONNX("models/best.onnx")
        self.pub = rospy.Publisher('/traffic_light', TrafficLightMsg, queue_size=10)

    def image_callback(self, msg):
        img = self.bridge.imgmsg_to_cv2(msg)
        blob = cv2.dnn.blobFromImage(img, 1/255.0, (896,896), swapRB=True)
        self.model.setInput(blob)
        outputs = self.model.forward()
        # 后处理逻辑...
        self.publish_detection(results)

在真实路测中发现，模型对黄昏时段的黄灯识别准确率会下降约15%。通过针对性采集500组黄昏数据做增量训练后，准确率回升到92%。这提醒我们： 数据分布的完整性比数据量更重要 。