YOLOv5 中的模型剪枝与稀疏性#
📚 本指南介绍如何对 YOLOv5 🚀 模型应用剪枝，在保持性能的同时创建更高效的网络。
什么是模型剪枝？#
模型剪枝是一种通过移除不太重要的参数（权重和连接）来缩减神经网络规模和复杂度的技术。此过程能创建更高效的模型，并带来以下优势：
- 缩小模型体积，更便于在资源受限的设备上部署
- 在对准确率影响极小的情况下加快推理速度
- 降低内存使用量和能耗
- 提高实时应用的整体效率
剪枝通过识别并移除对模型性能贡献很小的参数，使模型变得更轻量，同时保持相近的准确率。
开始之前#
克隆仓库，并在包含 PyTorch>=1.8 的 Python>=3.8.0 环境中安装 requirements.txt。模型和数据集会从 YOLOv5 最新版本自动下载。
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install测试基准性能#
剪枝之前，先确定基准性能以便进行比较。此命令以 640 像素的图像尺寸在 COCO val2017 上测试 YOLOv5x。yolov5x.pt 是可用的最大、精度最高的模型。其他选项包括 yolov5s.pt、yolov5m.pt 和 yolov5l.pt，也可以使用你在自定义数据集上训练得到的检查点，例如 runs/train/exp/weights/best.pt。如需了解所有可用模型的详细信息，请参阅 README 表格。
python val.py --weights yolov5x.pt --data coco.yaml --img 640 --half输出：
val: data=/content/yolov5/data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True, dnn=False
YOLOv5 🚀 v6.0-224-g4c40933 torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)
Fusing layers...
Model Summary: 444 layers, 86705005 parameters, 0 gradients
val: Scanning '/content/datasets/coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupt: 100% 5000/5000 [00:00<?, ?it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 157/157 [01:12<00:00, 2.16it/s]
all 5000 36335 0.732 0.628 0.683 0.496
Speed: 0.1ms pre-process, 5.2ms inference, 1.7ms NMS per image at shape (32, 3, 640, 640) # <--- base speed
Evaluating pycocotools mAP... saving runs/val/exp2/yolov5x_predictions.json...
...
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.507 # <--- base mAP
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.689
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.552
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.345
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.559
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.652
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.381
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.630
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.682
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.526
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.731
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.829
Results saved to runs/val/exp2对 YOLOv5x 应用剪枝（稀疏度为 30%）#
我们可以使用 utils/torch_utils.py 中定义的 prune() 函数对模型进行剪枝。要测试剪枝后的模型，我们将 val.py 更新为对 YOLOv5x 进行剪枝，使其稀疏度达到 0.3（30% 的权重设为零）：
剪枝 30% 后的输出：
val: data=/content/yolov5/data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True, dnn=False
YOLOv5 🚀 v6.0-224-g4c40933 torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)
Fusing layers...
Model Summary: 444 layers, 86705005 parameters, 0 gradients
Pruning model... 0.3 global sparsity
val: Scanning '/content/datasets/coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupt: 100% 5000/5000 [00:00<?, ?it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 157/157 [01:11<00:00, 2.19it/s]
all 5000 36335 0.724 0.614 0.671 0.478
Speed: 0.1ms pre-process, 5.2ms inference, 1.7ms NMS per image at shape (32, 3, 640, 640) # <--- prune speed
Evaluating pycocotools mAP... saving runs/val/exp3/yolov5x_predictions.json...
...
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.489 # <--- prune mAP
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.677
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.537
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.334
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.542
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.635
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.370
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.612
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.664
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.496
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.722
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.803
Results saved to runs/val/exp3结果分析#
从结果中可以观察到：
- 稀疏度达到 30%：
nn.Conv2d层中模型权重参数的 30% 现在为零 - 推理时间保持不变：尽管进行了剪枝，处理速度基本没有变化
- 性能影响极小：mAP 从 0.507 略微下降到 0.489（仅降低 3.6%）
- 模型大小未变：置零的权重仍以稠密形式存储，因此只有采用支持稀疏的格式或压缩后，文件大小才会缩小
这表明，剪枝只会轻微影响性能，却能显著降低模型复杂度，因此是面向资源受限环境进行部署时的有效优化技术。
微调剪枝后的模型#
为了获得最佳效果，应在剪枝后对模型进行微调，以恢复准确率。可以按以下步骤进行：
- 按照所需的稀疏度进行剪枝
- 使用较低的学习率对剪枝后的模型训练几个周期
- 将微调后的剪枝模型与基准模型进行评估对比
这个过程有助于剩余参数进行调整，以弥补被移除的连接，通常可以恢复大部分或全部原始准确率。
支持的环境#
Ultralytics 提供多种开箱即用的环境，每种环境都预装了 CUDA、CUDNN、Python 和 PyTorch 等必备依赖，助你快速启动项目。
- 免费 GPU 笔记本：
- Google Cloud：GCP 快速入门指南
- Amazon：AWS 快速入门指南
- Azure：AzureML 快速入门指南
- Docker：Docker 快速入门指南
项目状态#
此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些 CI 测试会严格检查 YOLOv5 在多个关键方面的功能和性能：训练、验证、推理、导出和基准测试。这些测试每 24 小时以及每次提交新代码后运行一次，确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上保持稳定可靠地运行。
常见问题#
非结构化剪枝会将权重设为零，但不会改变张量形状，因此稠密 PyTorch 内核仍会执行相同数量的运算。要获得加速效果，需要使用支持稀疏性的运行时（例如 DeepSparse），或采用会移除整个通道的结构化剪枝。
在验证或导出之前，调用
utils/torch_utils.py中的辅助函数：from utils.torch_utils import prune prune(model, amount=0.3) # 对 nn.Conv2d 权重设置 30% 的全局稀疏度从 0.3 开始并测量 mAP。YOLOv5x 在稀疏度为 30% 且未进行微调时，mAP 下降了约 3.6%；更高的稀疏度通常需要以较低的学习率重新训练几个周期，才能恢复准确率。
可以。置零的权重会像其他权重一样保存，因此
export.py无需更改即可正常工作。只有目标格式或运行时会压缩稀疏张量时，文件大小才会缩小。