YOLOv5模型剪枝压缩

基于yolov5 v5.0分支进行剪枝，采用yolov5s模型，原理为Learning Efficient Convolutional Networks Through Network Slimming（https://arxiv.org/abs/1708.06519）。

yolov5 v5.0转NCNN和安卓移植见https://blog.csdn.net/IEEE_FELLOW/article/details/117190219

yolov5s是非常优秀的轻量级检测网络，但是有时候模型依然比较大，使得我们不得不缩减网络输入大小，但是单纯降低输入来减少运算，例如640降低到320，对检测效果损失很大，同时模型体积依然是14M左右，所以可以通过添加L1正则来约束BN层系数，使得系数稀疏化，通过稀疏训练后，裁剪掉稀疏很小的层，对应激活也很小，所以对后面的影响非常小，反复迭代这个过程，可以获得很compact的模型，步骤基本是这样。

让我们扩展一下原理：

我们知道BN层的计算是这样的：

所以每个channel激活大小Zout和系数γ（pytorch对应bn层的weights，β对应bias）正相关，如果γ太小接近于0，那么激活值也非常小：

那么拿掉那些γ->0的channel是可以的，但是正常情况下，我们训练一个网络后，bn层的系数是类似正态分布：

上图就是正常训练时候γ的随着epoch的直方图分布，可以看基本正太分布。0附近的值是很少的，所以没法剪枝。

通过添加L1 正则约束：

上面第一项是正常训练的loss函数，第二项是约束，其中g(s) = |s|，λ是正则系数，根据数据集调整。可以将参数稀疏化，看看如果添加到训练的损失函数中去，在进行反向传播时候：

𝐿′=∑𝑙′+𝜆∑𝑔′(𝛾)=∑𝑙′+𝜆∑|𝛾|′=∑𝑙′+𝜆∑𝛾∗𝑠𝑖𝑔𝑛(𝛾)

所以，只需要在训练中，反向传播时候，在BN层权重乘以权重的符合函数输出和系数即可，对应添加如下代码:

另请注意以下事项：

1. yolov5会采用自动混合精度训练，这里改成fp32模式；

2。yolov5里面会通过model.fuse()将卷积层和bn层融合，为了对bn层剪枝，训练和保存时候先不要fuse。

3。这里并没有选择所有的bn层进行裁剪，这里选择去除那些有shortcut的Bottleneck层(对应代码中m.add = True)，主要是为了保证shortcut和残差层channel一样可以add。

4.稀疏训练后，将BN层前的卷积层对应通道的卷积核裁剪掉，将BN层后对应的特征图裁剪掉。

ok，我们选用一个较小的22类数据集训练尝试，λ从0.001均匀变化到0.0001训练完成，看看bn层变化：

可以明显看到，随着训练进行（纵轴是epoch），BN层参数逐渐从最上面的正太分布趋向于0附近。

这里裁剪60%的参数是试试：

修剪前：

修剪后：

可以看到剪枝60%的通道后，MAP仅仅降低0.5%，通过微调可以很快恢复精度。

频道剪辑：

上图是裁剪后的模型的配置，这里重写了Bottleneck,C3,SPP三个模块。

剪枝前后模型体积对比：

反复迭代这个过程，就可以获得一个更加轻量的yolov5s模型。

关于yolov5s模型分析和剪枝选择分析见：https://blog.csdn.net/IEEE_FELLOW/article/details/117536808

所有源码和实验见：https://github.com/midasklr/yolov5prune