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层权重乘以权重的符合函数输出和系数即可,对应添加如下代码:
同时需要注意以下几点:
- yolov5会采用自动混合精度训练,这里改成fp32模式;
- yolov5里面会通过model.fuse()将卷积层和bn层融合,为了对bn层剪枝,训练和保存时候先不要fuse。
- 这里并没有选择所有的bn层进行裁剪,这里选择去除那些有shortcut的Bottleneck层(对应代码中m.add = True),主要是为了保证shortcut和残差层channel一样可以add。
- 稀疏训练后,将BN层前的卷积层对应通道的卷积核裁剪掉,将BN层后对应的特征图裁剪掉。
选用一个较小的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