YOLOv5中的模型剪枝和稀疏性
本指南介绍如何对YOLOv5 🚀 模型进行剪枝,以创建更高效的网络,同时保持性能。
什么是模型修剪?
模型剪枝是一种通过删除不那么重要的参数(权重和连接)来减小神经网络规模和复杂性的技术。这一过程可以创建一个更高效的模型,并带来多种益处:
- 缩小模型尺寸,更易于在资源有限的设备上部署
- 推理速度更快,对准确性的影响最小
- 降低内存使用率和能耗
- 提高实时应用的整体效率
剪枝的工作原理是识别并移除对模型性能贡献最小的参数,从而得到一个更轻量级、精度相似的模型。
开始之前
克隆 repo 并将requirements.txt安装在 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
测试基准性能
在剪枝之前,先建立一个基准性能来进行比较。此命令在 COCO val2017 上测试 YOLOv5x,图像大小为 640 像素。 yolov5x.pt
是目前最大、最精确的模型。其他选项包括 yolov5s.pt
, yolov5m.pt
和 yolov5l.pt
或您自己从训练自定义数据集中获得的检查点 ./weights/best.pt
.有关所有可用型号的详细信息,请参见 README 桌.
输出:
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_hybrid=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/exp
对 YOLOv5x 进行修剪(30% 稀疏度)
我们可以使用 torch_utils.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_hybrid=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% 的稀疏性:模型权重参数的 30%。
nn.Conv2d
层数现在为零 - 推理时间保持不变:尽管进行了剪枝处理,但处理速度基本不变
- 对性能的影响极小:mAP 从 0.507 微降至 0.489(仅降低 3.6)
- 缩小模型尺寸:剪枝模型所需的存储内存更少
这表明,剪枝可以大大降低模型的复杂性,而对性能的影响很小,因此是在资源受限环境中部署的一种有效优化技术。
微调剪枝模型
为获得最佳效果,剪枝后应对剪枝模型进行微调,以恢复准确性。具体做法如下
- 根据所需的稀疏程度进行剪枝
- 用较低的学习率训练剪枝模型几个历元
- 对照基线评估经过微调的剪枝模型
这一过程有助于其余参数进行调整,以补偿被移除的连接,通常能恢复大部分或全部原始精度。
支持的环境
Ultralytics 提供了一系列随时可用的环境,每个环境都预装了基本的依赖项,如 CUDA、CUDNN、 Python和 PyTorch等基本依赖项,以便启动项目。
- 免费GPU 笔记本:
- Google 云计算 GCP 快速入门指南
- 亚马逊 AWS 快速入门指南
- Azure.AzureML 快速入门指南AzureML 快速入门指南
- Docker: Docker 快速入门指南
项目现状
此徽章表示YOLOv5 GitHub Actions 的所有持续集成(CI)测试均已成功通过。这些 CI 测试严格检查了YOLOv5 在训练、验证、推理、导出和基准等多个关键方面的功能和性能。它们确保在 macOS、Windows 和 Ubuntu 上运行的一致性和可靠性,每 24 小时和每次新提交时都会进行一次测试。