yolov5模型修剪/稀疏性
在你开始前
克隆 YOLOv5 存储库并安装requirements.txt依赖项,包括Python>=3.8和PyTorch>=1.7。
$ git clone https://github.com/ultralytics/yolov5  # clone repo
$ cd yolov5
$ pip install -r requirements.txt  # install
在 COCO 上测试 YOLOv5x(默认)
此命令在 COCO val2017 上以 640 像素的图像大小测试 YOLOv5x 以建立标称基线。yolov5x.pt是可用的最大和最准确的模型。其他选项是yolov5s.pt, yolov5m.ptand yolov5l.pt, 或者您拥有训练自定义数据集的检查点./weights/best.pt。有关所有可用模型的详细信息,请参阅我们的自述文件表。
$ python test.py --weights yolov5x.pt --data coco.yaml --img 640 --iou 0.65
默认输出:
YOLOv5 🚀 v4.0-174-g9c803f2 torch 1.8.0+cu101 CUDA:0 (Tesla V100-SXM2-16GB, 16160.5MB)
Fusing layers... 
Model Summary: 476 layers, 87730285 parameters, 0 gradients, 218.8 GFLOPS
val: Scanning '../coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100% 5000/5000 [00:00<00:00, 50901747.57it/s]
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100% 157/157 [01:20<00:00,  1.95it/s]
                 all        5000       36335       0.749       0.619        0.68       0.486
Speed: 5.2/1.6/6.8 ms inference/NMS/total per 640x640 image at batch-size 32  < -------- speed
Evaluating pycocotools mAP... saving runs/test/exp/yolov5x_predictions.json...
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.501  < -------- mAP
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.687
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.544
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.338
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.548
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.637
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.378   < -------- mAR
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.628
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.680
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.520
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.729
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.826
Results saved to runs/test/exp
在 COCO 上测试 YOLOv5x(0.30 稀疏度)
我们使用torch_utils.prune()命令对修剪后的模型重复上述测试。我们更新test.py以将 YOLOv5x 修剪为 0.3 稀疏度:

30% 修剪输出:
YOLOv5 🚀 v4.0-174-g9c803f2 torch 1.8.0+cu101 CUDA:0 (Tesla V100-SXM2-16GB, 16160.5MB)
Fusing layers... 
Model Summary: 476 layers, 87730285 parameters, 0 gradients, 218.8 GFLOPS
Pruning model...  0.3 global sparsity
val: Scanning '../coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupted: 100% 5000/5000 [00:00<00:00, 50901747.57it/s]
               Class      Images      Labels           P           R      mAP@.5  mAP@.5:.95: 100% 157/157 [01:17<00:00,  2.01it/s]
                 all        5000       36335       0.742       0.595        0.66        0.46
Speed: 5.2/1.5/6.7 ms inference/NMS/total per 640x640 image at batch-size 32  < -------- speed
Evaluating pycocotools mAP... saving runs/test/exp/yolov5x_predictions.json...
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.477  < -------- mAP
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.667
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.525
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.313
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.530
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.619
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.366   < -------- mAR
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.603
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.654
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.475
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.709
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.794
Results saved to runs/test/exp
在结果中我们可以观察到,我们的模型在剪枝后达到了30%的稀疏性,这意味着模型nn.Conv2d层中30% 的权重参数等于 0。推理时间基本不变,而模型的AP 和 AR分数略有降低。
评论
