YOLOv5 中的模型剪枝与稀疏性#
📚 本指南介绍如何对 剪枝 YOLOv5 🚀 模型，以创建更高效的网络，同时保持性能。
什么是模型剪枝？#
模型剪枝 是一种通过移除不太重要的参数（权重和连接）来减小神经网络规模和复杂度的技术。此过程可以创建更高效的模型，并带来以下优势：
- 减小模型大小，更易于部署到资源受限的设备上
- 在对准确率影响极小的情况下提升推理速度
- 降低内存使用量和能耗
- 提升实时应用的整体效率
剪枝通过识别并移除对模型性能贡献很小的参数来实现，从而得到一个更轻量的模型，同时保持相近的准确率。
开始之前#
在 Python>=3.8.0 环境中克隆仓库并安装 requirements.txt，其中包括 PyTorch>=1.8。模型 和数据集 会自动从最新的 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，或者你也可以使用在自定义数据集上训练得到的检查点 ./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% 稀疏率）#
我们可以使用 torch_utils.prune() 命令对模型进行剪枝，该命令定义于 utils/torch_utils.py 中。要测试剪枝后的模型，我们将 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 在多个关键方面的功能和性能：训练、验证、推理、导出和基准测试。它们确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上持续稳定地运行，并且每 24 小时以及每次新提交时都会执行测试。