YOLOv5 v6.0 添加 P2 检测头实战：VisDrone 数据集 mAP 提升 2.9%

YOLOv5 v6.0 添加 P2 检测头实战：VisDrone 数据集 mAP 提升 2.9%

无人机航拍场景下的目标检测一直是计算机视觉领域的难点，尤其是对于小目标的识别。VisDrone 数据集作为无人机视角下的标准 benchmark，其密集小目标的特性对检测算法提出了更高要求。本文将手把手带您完成 YOLOv5 v6.0 添加 P2 检测头的完整流程，并通过实验验证其在 VisDrone 数据集上的效果提升。

1. 检测头原理与 P2 层的价值

目标检测模型通常由三部分组成：Backbone（主干网络）、Neck（特征融合层）和 Head（检测头）。检测头作为最终预测的关键组件，其设计直接影响模型对不同尺度目标的敏感度。

传统 YOLOv5 采用 P3、P4、P5 三级检测头，分别对应不同尺度的特征图：

对于 VisDrone 这类无人机数据集，目标尺寸分布呈现明显特点：

# VisDrone 目标尺寸分布统计（像素） small_objects = 62% # <32×32 medium_objects = 28% # 32×32~96×96 large_objects = 10% # >96×96

P2 检测头（4x 下采样，160×160 特征图）的引入带来三个核心优势：

1. 更早保留小目标的细节特征
2. 减少深层卷积对小目标特征的"淹没效应"
3. 提升密集场景下的目标区分能力

注意：P2 层虽然提升小目标检测，但会增加约15%的计算量，需权衡精度与速度

2. 配置文件修改实战

YOLOv5 的模块化设计使得添加 P2 检测头仅需修改 YAML 配置文件，无需改动代码。以下是完整的yolov5-p2.yaml配置：

# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license # Parameters nc: 10 # VisDrone 类别数 depth_multiple: 0.33 width_multiple: 0.50 anchors: 4 # 调整为4个检测头 # Backbone backbone: # [from, number, module, args] [[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2 [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 [-1, 3, C3, [128]], [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 [-1, 6, C3, [256]], [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 [-1, 9, C3, [512]], [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 [-1, 3, C3, [1024]], [-1, 1, SPPF, [1024, 5]], # 9 ] # Head with P2-P5 outputs head: [[-1, 1, Conv, [512, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], [[-1, 6], 1, Concat, [1]], # cat P4 [-1, 3, C3, [512, False]], # 13 [-1, 1, Conv, [256, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], [[-1, 4], 1, Concat, [1]], # cat P3 [-1, 3, C3, [256, False]], # 17 [-1, 1, Conv, [128, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], [[-1, 2], 1, Concat, [1]], # cat P2 [-1, 3, C3, [128, False]], # 21 [-1, 1, Conv, [128, 3, 2]], [[-1, 18], 1, Concat, [1]], # cat P3 [-1, 3, C3, [256, False]], # 24 [-1, 1, Conv, [256, 3, 2]], [[-1, 14], 1, Concat, [1]], # cat P4 [-1, 3, C3, [512, False]], # 27 [-1, 1, Conv, [512, 3, 2]], [[-1, 10], 1, Concat, [1]], # cat P5 [-1, 3, C3, [1024, False]], # 30 [[21, 24, 27, 30], 1, Detect, [nc, anchors]], # Detect(P2,P3,P4,P5) ]

关键修改点说明：

1. anchors: 4调整为4个检测头
2. Backbone 中明确保留 P2 层特征（第1层输出）
3. Head 部分增加 P2 分支（第21层）
4. Detect 层接收四个检测头的输入

3. 训练与验证

使用 VisDrone 数据集进行训练时，建议采用以下参数组合：

python train.py \ --data visdrone.yaml \ --cfg yolov5-p2.yaml \ --weights yolov5s.pt \ --batch-size 16 \ --img 1024 \ --epochs 100 \ --hyp hyp.visdrone.yaml \ --name yolov5s-p2

VisDrone 专用超参数配置（hyp.visdrone.yaml）：

# Hyperparameters for small object detection 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.301 obj_pw: 0.911 iou_t: 0.2 anchor_t: 3.0 fl_gamma: 1.5 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.243 copy_paste: 0.0

训练完成后，使用验证脚本测试性能：

python val.py \ --data visdrone.yaml \ --weights runs/train/yolov5s-p2/weights/best.pt \ --img 1024 \ --task test \ --save-json

4. 实验结果对比

我们在 VisDrone 测试集上对比了三种配置：

小目标检测提升明显（<32px 目标）：

可视化对比显示，P2 检测头在密集小目标场景中表现优异：

5. 工程优化建议

在实际部署中，我们总结了以下优化经验：

1. Anchor 调整：

   # VisDrone 专用 anchor 配置（使用 k-means 重新聚类） anchors: - [4,5, 8,10, 13,16] # P2 - [19,24, 38,48, 57,71] # P3 - [86,107, 142,177, 273,341] # P4

2. 混合精度训练：

   python train.py --device 0 --batch-size 32 --img 1024 --adam --half

3. 模型剪枝（适用于边缘设备）：

   # 示例剪枝代码（需安装 torch_pruner） from torch_pruner import prune_model prune_model(model, method='l1', amount=0.3, exclude=['Detect'])

4. TensorRT 加速：

   python export.py --weights best.pt --include engine --device 0 --half

经过实际项目验证，在 Jetson Xavier NX 设备上，优化后的 P234 模型能保持 25 FPS 的实时性能，满足大多数无人机应用场景需求。