YOLOv5 压缩剪枝

github： midasklr/yolov5prune at v6.0 (github.com)

giithub作者讲解：(6条消息) yolov5模型压缩之模型剪枝_yolov5模型剪枝_小小小绿叶的博客-CSDN博客

yolov5剪枝实操：(7条消息) YOLOv5模型剪枝压缩_yolov5剪枝_MidasKing的博客-CSDN博客

(7条消息) yolov5s模型剪枝详细过程(v6.0)_yolov5s剪枝_银嘉诚的博客-CSDN博客

本文参考上面的博客，复现yolov5 6.0中的BN剪枝操作。记录下当时的复现流程。

BN层剪枝（yolov5 6.0）

剪枝流程：1.训练出一个最优权重；2.使用最优权重进行稀疏化训练；3.按比例裁剪；4.微调

这里使用了自己搭建的私人数据集，大家可以选择自己需要的数据集，修改相关yaml中的参数信息即可。yolov5输入大小为512.为了节省时间,这里所有实验使用AdamW训练100 epoch.使用SGD+300epoch+0.01=>0.001的参数设置通常会有更好的结果.

1.正常训练:

python train.py --weights yolov5s.pt --data data/coco128.yaml --epochs 10 --imgsz 512 --adam 

首先使用yolov5s.pt训练出一个基础权重。训练数据集选用私人小数据集，浅训10epoch，相应修改yolov5s.yaml 的 nc：3；取得预训练baseline；

2.稀疏化训练

python train_sparity.py --st --sr 0.0002 --weights yolov5s.pt --data data/voc.yaml --epochs 100 --imgsz 512 --adam ...

训练中需要通过tensorboard监控训练过程,特别是map变化,bn分布变化等,在runs/train/exp*/目录下有events.out.tfevents.* 文件,在此目录下输入:

tensorboard --logdir .

然后点击出现的链接观察训练中的各项指标.

• tensorboard可视化时报错：Duplicate plugins for name projector

• 解决：报错信息的直观含义是重复插件(安装了多个tensorboard)。
  不需要删除，改变环境。
  只需要到 Anaconda3\Lib\site-packages(tensorboard所对应的虚拟环境下) 把tensorboard-2.2.2.dist-info 删除即可

正常训练的bn参数γ处于正态分布，通过添加L1正则约束，将γ均值趋近于0，方便之后剪枝，对应代码就是下图在loss反向传播时候，在BN层权重乘以权重的符合函数输出和系数即可，

1. yolov5会采用自动混合精度训练，这里改成fp32模式；
2. yolov5里面会通过model.fuse()将卷积层和bn层融合，为了对bn层剪枝，训练和保存时候先不要fuse。
3. 这里并没有选择所有的bn层进行裁剪，这里选择去除那些有shortcut的Bottleneck层(对应代码中m.add = True)，主要是为了保证shortcut和残差层channel一样可以add。
4.稀疏训练后，将BN层前的卷积层对应通道的卷积核裁剪掉，将BN层后对应的特征图裁剪掉。

使用tensorboard可视化，可以看到参数γ逐渐趋近于0，下一步开始剪枝

3.对bn层进行剪枝

python prune.py --weights last.pt --percent 0.1 --data data/coco128.yaml

裁剪比例percent根据效果调整。注意cfg的模型文件需要和weights对应上,否则会出现运行prune 过程中出现键值不对应的问题,裁剪完成会保存对应的模型pruned_model.pt。裁剪时需要注意channel需要保持为4的倍数，不然可能导致裁剪后推理速度减慢。

找到prune_utils.py 将函数obtain_bn_mask（）替换掉下图中的函数。

运行过程中会看到下面的提示：

Suggested Gamma threshold should be less than 0.6030.

The corresponding prune ratio is 0.264, but you can set higher.

Gamma value that less than 0.7095 are set to zero!

这里剪枝的多少取决于percent参数，裁剪中会出现裁剪掉的gamma值提示，percent选择0.6时，这里小于0.7095的γ值被裁减掉。裁剪后进行微调，下面是原始指标和微调后指标对比.

Class Images Labels P R [email protected] [email protected]:.95

all 168 611 0.931 0.912 0.946 0.68

person 168 236 0.945 0.89 0.954 0.647

head 168 223 0.973 0.971 0.985 0.656

screen 168 152 0.875 0.875 0.899 0.738

Class Images Labels P R [email protected] [email protected]:.95

all 168 611 0.927 0.934 0.944 0.707

person 168 236 0.971 0.941 0.966 0.704

head 168 223 0.96 0.969 0.976 0.654

screen 168 152 0.85 0.894 0.89 0.761

4.微调：

python finetune_pruned.py --weights pruned_model.pt --adam --epochs 100

在数据集上的实验结果:

	模型大小		cpu	3090
yolov5s			Speed: 0.9ms pre-process, 302.5ms inference, 0.7ms NMS per image at shape (1, 3, 640, 640)
yolov5s	14M	213 layers, 7018216 parameters, 0 gradients, 15.8 GFLOPs	Speed: 0.6ms pre-process, 196.3ms inference, 0.6ms NMS per image at shape (1, 3, 512, 512)	Speed: 0.4ms pre-process, 3.2ms inference, 8.2ms NMS per image at shape (32, 3, 512, 512)
prune20%yolov5s	10.4M		Speed: 0.6ms pre-process, 189.8ms inference, 0.4ms NMS per image at shape (1, 3, 512, 512)
prune60%yolov5s	5.66M		Speed: 0.7ms pre-process, 109.5ms inference, 0.6ms NMS per image at shape (1, 3, 512, 512)	Speed: 0.2ms pre-process, 2.7ms inference, 15.3ms NMS per image at shape (32, 3, 512, 512)