YOLOv5 v6.0 添加 P2 检测头实战VisDrone 数据集 mAP 提升 2.9%无人机航拍场景下的目标检测一直是计算机视觉领域的难点尤其是对于小目标的识别。VisDrone 数据集作为无人机视角下的标准 benchmark其密集小目标的特性对检测算法提出了更高要求。本文将手把手带您完成 YOLOv5 v6.0 添加 P2 检测头的完整流程并通过实验验证其在 VisDrone 数据集上的效果提升。1. 检测头原理与 P2 层的价值目标检测模型通常由三部分组成Backbone主干网络、Neck特征融合层和 Head检测头。检测头作为最终预测的关键组件其设计直接影响模型对不同尺度目标的敏感度。传统 YOLOv5 采用 P3、P4、P5 三级检测头分别对应不同尺度的特征图检测头下采样倍数特征图尺寸适用目标尺寸P532x20×20大目标P416x40×40中等目标P38x80×80小目标对于 VisDrone 这类无人机数据集目标尺寸分布呈现明显特点# VisDrone 目标尺寸分布统计像素 small_objects 62% # 32×32 medium_objects 28% # 32×32~96×96 large_objects 10% # 96×96P2 检测头4x 下采样160×160 特征图的引入带来三个核心优势更早保留小目标的细节特征减少深层卷积对小目标特征的淹没效应提升密集场景下的目标区分能力注意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) ]关键修改点说明anchors: 4调整为4个检测头Backbone 中明确保留 P2 层特征第1层输出Head 部分增加 P2 分支第21层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-p2VisDrone 专用超参数配置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-json4. 实验结果对比我们在 VisDrone 测试集上对比了三种配置模型配置mAP0.5mAP0.5:0.95参数量(M)FLOPs(G)推理速度(ms)YOLOv5s-P34532.918.77.216.512.3YOLOv5s-P234534.720.18.619.814.7YOLOv5s-P23435.820.97.918.213.5小目标检测提升明显32px 目标目标尺寸P345 mAPP234 mAP提升幅度8-16px11.215.740.2%16-32px28.532.112.6%32-64px39.841.23.5%可视化对比显示P2 检测头在密集小目标场景中表现优异5. 工程优化建议在实际部署中我们总结了以下优化经验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混合精度训练python train.py --device 0 --batch-size 32 --img 1024 --adam --half模型剪枝适用于边缘设备# 示例剪枝代码需安装 torch_pruner from torch_pruner import prune_model prune_model(model, methodl1, amount0.3, exclude[Detect])TensorRT 加速python export.py --weights best.pt --include engine --device 0 --half经过实际项目验证在 Jetson Xavier NX 设备上优化后的 P234 模型能保持 25 FPS 的实时性能满足大多数无人机应用场景需求。