基于YOLOv5的骨干网络替换与模型剪枝优化（支持TensorRT与NCNN部署）
YOLOv5 改进概览
本文介绍对YOLOv5系列目标检测模型的一系列结构改进，包括主干网络（Backbone）替换、轻量化设计、模型剪枝及量化训练，并支持在边缘设备上高效部署。所有实验均基于 VisDrone 数据集进行验证，输入分辨率涵盖608至1536。
- 2022.10.30：实现TPH-YOLOv5结构复现
- 2022.10.31：集成GhostNet作为主干网络
- 2022.11.02：接入ShuffleNetV2主干
- 2022.11.05：引入MobileNetV3-Small主干
- 2022.11.10：支持EagleEye通道剪枝算法
- 2022.11.14：集成MQBench量化训练框架
- 2022.11.16：采用EfficientNet-Lite0为主干
- 2022.11.26：融合PP-LCNet-1x主干
- 2022.12.12：构建SwinTransformer增强型C3STR模块
- 2022.12.15：完成Network Slimming剪枝支持
环境依赖
pip install -r requirements.txt
多主干网络替换方案
基准模型性能对比（VisDrone数据集，输入尺寸608）
| No. | Model | mAP | mAP@50 | Parameters(M) | GFLOPs |
|---|---|---|---|---|---|
| 1 | YOLOv5n | 13.0 | 26.20 | 1.78 | 4.2 |
| 2 | YOLOv5s | 18.4 | 34.00 | 7.05 | 15.9 |
| 3 | YOLOv5m | 21.6 | 37.80 | 20.91 | 48.2 |
| 4 | YOLOv5l | 23.2 | 39.70 | 46.19 | 108.1 |
| 5 | YOLOv5x | 24.3 | 40.80 | 86.28 | 204.4 |
高精度模型变体
TPH-YOLOv5
通过引入P2输出层、CBAM注意力机制、TPH（Transformer Prediction Head）、BiFPN特征融合和SPP模块，显著提升小目标检测能力。
| No. | Model | mAP | mAP@50 | Params(M) | GFLOPs |
|---|---|---|---|---|---|
| 6 | YOLOv5xP2 | 30.0 | 49.29 | 90.96 | 314.2 |
| 7 | YOLOv5xP2 + CBAM | 30.1 | 49.40 | 91.31 | 315.1 |
| 8 | YOLOv5x-TPH | 40.7 | 63.00 | 112.97 | 270.8 |
训练命令示例
nohup python train.py --data VisDrone.yaml --weights yolov5n.pt --cfg models/yolov5n.yaml --epochs 300 --batch-size 8 --img 608 --device 0,1 --sync-bn >> yolov5n.txt &
关键设计细节
- TransBlock数量随YOLO版本调整，标准结构适用于YOLOv5m
- 对于YOLOv5x，移除第14和19层的CBAM模块，降低P2路径通道数至128
- SPP模块置于输出头前，其卷积核大小随空间分辨率动态缩放
- 仅保留backbone末端和最终预测层的TransBlock
- Neck部分采用BiFPN结构增强多尺度融合
损失函数权重调优
当验证集分类与置信度损失出现上升趋势时，表明过拟合发生，需降低对应任务权重：
| box | cls | obj | mAP |
|---|---|---|---|
| 0.05 | 0.5 | 1.0 | 37.90 |
| 0.05 | 0.3 | 0.7 | 38.00 |
| 0.05 | 0.2 | 0.4 | 37.50 |
SwinTrans-YOLOv5
结合Swin Transformer的局部窗口自注意力机制，提升长距离依赖建模能力。
安装依赖
pip install timm
训练命令
python train.py --data VisDrone.yaml --weights yolov5x.pt --cfg models/accModels/yolov5xP2CBAM-Swin-BiFPN-SPP.yaml --hyp data/hyps/hyp.visdrone.yaml --epochs 60 --batch-size 4 --img 1536 --nohalf
实现要点
- 将原始window size从7改为8，适配常见图像分辨率
- 将mask生成逻辑由
__init__迁移至forward中执行 - 若输入尺寸小于window size或非整倍数，则在右下角补零
- 计算完成后需去除padding以保持与其他分支对齐
- 前后需对张量进行reshape操作以匹配Transformer输入格式
- 使用C3STR模块时需禁用半精度验证（添加
--nohalf参数）
轻量级模型比较（VisDrone数据集）
| No | Model | mAP | mAP@50 | Params(M) | GFLOPs | TrainCost(h) | Memory(G) | FPS@CPU |
|---|---|---|---|---|---|---|---|---|
| 1 | YOLOv5l | 23.2 | 39.7 | 46.19 | 108.1 | - | - | - |
| 2 | YOLOv5l-GhostNet | 18.4 | 33.8 | 24.27 | 42.4 | 27.44 | 4.97 | - |
| 3 | YOLOv5l-ShuffleNetV2 | 16.48 | 31.1 | 21.27 | 40.5 | 10.98 | 2.41 | - |
| 4 | YOLOv5l-MobileNetv3Small | 16.55 | 31.2 | 20.38 | 38.4 | 10.19 | 5.30 | - |
| 5 | YOLOv5l-EfficientNetLite0 | 19.12 | 35.0 | 23.01 | 43.9 | 13.94 | 2.04 | - |
| 6 | YOLOv5l-PP-LCNet | 17.63 | 32.8 | 21.64 | 41.7 | 18.52 | 1.66 | - |
各轻量主干适配策略
GhostNet-YOLOv5
- 统一使用kernel_size=3的深度可分离卷积进行下采样
- Neck与Head结构沿用YOLOv5l原版
- 中间扩展通道数人为设定（expand ratio控制）
ShuffleNetV2-YOLOv5
- 摒弃Focus层，避免slice操作导致缓存效率低下
- 减少C3模块重复堆叠次数，尤其避免高通道数下的密集使用
- 通道数保持不变以符合组卷积设计原则
MobileNetV3Small-YOLOv5
- 保留Hard-Swish激活函数和SE注意力模块
- Neck与Head部分直接复用YOLOv5l结构
- 中间通道按预设比例扩展
EfficientNetLite0-YOLOv5
- 采用Lite版本结构，去除SE模块以提升推理速度
- dropout_connect_rate随stage递增逐步增大
- 所有中间通道统一乘以6作为扩展系数
PP-LCNet-YOLOv5
- 使用PP-LCNet-1x主干，在末端加入SE模块和5×5卷积核
- SE模块压缩比设置为1/16
- 通道维度整体保持不变
模型剪枝方法
剪枝后性能对比
| Model | mAP | mAP@50 | Params(M) | GFLOPs | FPS@CPU |
|---|---|---|---|---|---|
| YOLOv5s | 18.4 | 34.0 | 7.05 | 15.9 | - |
| YOLOv5n | 13.0 | 26.2 | 1.78 | 4.2 | - |
| YOLOv5s-EagleEye@0.6 | 14.3 | 27.9 | 4.59 | 9.6 | - |
剪枝策略设计
- 针对YOLOv5中的Conv、C3、SPP(F)模块进行通道剪枝
- 具体剪枝对象包括：
- Conv模块的输出通道
- C3中cv2和cv3的输出通道
- C3内bottleneck中cv1的输出通道
- 通道数约束为8的整数倍（out_channels = 8 × n）
- Shortcut连接与concat操作需同步剪枝以保证维度一致
剪枝工具链
EagleEye 剪枝流程
基于随机搜索的快速子网评估方法，通过BN层统计量校准筛选最优子结构。
- 训练原始模型
python train.py --data data/VisDrone.yaml --imgsz 640 --weights yolov5s.pt --cfg models/prunModels/yolov5s-pruning.yaml --device 0 - 搜索最优子网
python pruneEagleEye.py --weights path_to_trained_model --cfg models/prunModels/yolov5s-pruning.yaml --data data/VisDrone.yaml --path output_yaml_path --max_iter 1000 --remain_ratio 0.6 --delta 0.02 - 微调恢复精度
python train.py --data data/VisDrone.yaml --imgsz 640 --weights path_to_pruned_model --cfg path_to_pruned_yaml --device 0
Network Slimming 流程
利用BatchNorm层γ系数稀疏化驱动通道重要性学习。
- 稀疏训练
python train.py --data data/VisDrone.yaml --imgsz 640 --weights yolov5s.pt --cfg models/prunModels/yolov5s-pruning.yaml --device 0 --sparse - 执行剪枝
python pruneSlim.py --weights path_to_sparsed_model --cfg models/prunModels/yolov5s-pruning.yaml --data data/VisDrone.yaml --path output_yaml --global_percent 0.6 --device 3 - 微调恢复
python train.py --data data/VisDrone.yaml --imgsz 640 --weights path_to_Slimpruned_model --cfg path_to_pruned_yaml --device 0
量化感知训练（QAT）
基于MQBench框架实现硬件友好的量化训练，支持多种后端部署。
环境要求
- PyTorch == 1.8.1
安装MQBench
git clone https://github.com/ZLkanyo009/MQBench.git
cd MQBench
python setup.py build
python setup.py install
训练命令示例
python train.py --data VisDrone.yaml --weights yolov5n.pt --cfg models/yolov5n.yaml --epochs 300 --batch-size 8 --img 608 --nosave --device 0,1 --sync-bn --quantize --BackendType NNIE
部署支持
当前已支持以下推理引擎的模型导出与部署：
- TensorRT：适用于NVIDIA GPU，提供低延迟高性能推理
- NCNN：腾讯开源的移动端神经网络框架，无第三方依赖
详细部署流程请参考项目中的deploy目录说明。
