跳转至内容

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

测试模型在真实场景中的适应性和鲁棒性，超越常见的基准数据集，如 COCO 或 PASCAL VOC。

在剪枝之前，确定一个基准性能以便进行比较。此命令在 COCO val2017 数据集上以 640 像素的图像大小测试 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/exp

将剪枝应用于 YOLOv5x（30% 稀疏性）

我们可以使用以下方法对模型进行剪枝： torch_utils.prune() 命令。要测试修剪后的模型，我们更新 val.py 将 YOLOv5x 剪枝到 0.3 稀疏度（30% 的权重设置为零）：

显示将 YOLOv5x 剪枝到 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% 的稀疏度: 30%的模型权重参数在 nn.Conv2d 层现在为零
  • 推理时间保持不变: 尽管进行了剪枝，但处理速度基本相同
  • 性能影响极小: mAP 从 0.507 略微降至 0.489（仅降低 3.6%）。
  • 模型大小缩减: 剪枝后的模型需要更少的存储内存

这表明，剪枝可以显著降低模型复杂度，而对性能的影响很小，使其成为在资源受限环境中部署的有效优化技术。

微调剪枝模型

为了获得最佳效果，修剪后的模型应在修剪后进行微调以恢复准确性。 这可以通过以下方式完成：

  1. 应用具有所需稀疏度的剪枝
  2. 使用较低的学习率对剪枝后的模型进行几个 epoch 的训练
  3. 评估微调后的剪枝模型与基线的对比

此过程有助于剩余参数进行调整，以补偿移除的连接，通常可以恢复大部分或全部原始精度。

支持的环境

Ultralytics 提供一系列即用型环境，每个环境都预装了必要的依赖项，如 CUDA、CUDNN、Python 和 PyTorch，以快速启动您的项目。

项目状态

YOLOv5 CI

此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些 CI 测试严格检查 YOLOv5 在各个关键方面的功能和性能：训练、验证、推理、导出 和 基准测试。它们确保在 macOS、Windows 和 Ubuntu 上运行的一致性和可靠性，测试每 24 小时进行一次，并在每次提交新内容时进行。



📅 创建于 1 年前 ✏️ 更新于 2 个月前
glenn-jocher Laughing-q UltralyticsAssistant RizwanMunawar

评论