模型剪枝-yolov5

引言

为了了解学习模型剪枝,在此以目标检测算法yolov5为例子,进行模型模型剪枝操作。

一、概念介绍

1、模型的量化就是通过压缩或精简模型,减少模型的计算量和存储空间,从而提高模型的效率和速度。量化方式主要有两个方面:

(1)权重量化:将模型的浮点参数转换为低精度参数,比如32位浮点数转化为8位整数或者4位浮点数。

(2)模型剪枝:将模型不必要的参数和连接删除,减少计算量和内存使用空间,例如将不必要的卷积核删除,将稀疏连接剪枝掉。

二、Yolov5进行模型剪枝压缩

1、目标检测模型yolov5具有非常好的检测效果,但是如果模型体积较大,如果只是单纯降低网络输入来减少运算(如640->320),对检测效果损失会很大。通过添加L1正则化来约束BN层系数,使得系数稀疏化。通过稀疏训练后,裁剪掉稀疏层很小的层,对应的激活也很小,这样做对模型的性能的影响较低。反复迭代这个过程,可以获得很有效的模型,基本的步骤就是这样。

2、详细解释:BN层中缩放因子γ与卷积层中的每个通道关联起来。在训练过程中对这些比例因子进行稀疏正则化,以自动识别不重要的通道。缩放因子值较小的通道将被修剪。剪枝后,获得了紧凑的模型,然后对其进行微调,以达到与正常训练的全网络相当(甚至更高)的精度。

更进一步的解释:在BN层网络中加入稀疏因子,训练使得BN层稀疏化,对稀疏训练的后的模型中所有BN层权重进行统计排序,获取指定保留BN层数量即取得排序后权重阈值thres。遍历模型中的BN层权重,制作各层mask(权重>thres值为1,权重<thres值为0)。剪枝操作,根据各层的mask构建新模型结构(各层保留的通道数),获取BN层权重*mask非零值的索引,非零索引对应的原始conv层、BN层、linear层各通道的权重、偏置等值赋值给新模型各层。加载剪枝后模型,进行微调训练。

PS:正则化:

对原损失函数引入额外的信息,防止过拟合或者提高模型的泛化能力,也就是目标函数变成了原损失函数+额外项,常见的额外项就是L1正则化和L2正则化。

3、这里我们以行人手持危险物品作为训练集,yolov5prune作为训练代码(GitHub - midasklr/yolov5prune

(1)首先采用train.py进行正常的训练;

yolov5训练的完整模型  大小:13.6M 速度:0.9ms(413张图片)

(2)然后进行稀疏训练,train_sparsity.py;

yolov5经过稀疏训练得到的模型   大小:13.6M 速度:0.8ms(413张图片)

(3)训练完成后需要进行剪枝操作:prune.py;

模型:yolov5剪枝后得到的模型   大小:9.67M 速度0.8ms(413张图片)

(4)剪枝完成后,进行微调:finetune_pruned.py。

yolov5剪枝后重新微调的模型   大小:4.95M 速度0.8ms

猜你喜欢

转载自blog.csdn.net/qq_39149619/article/details/132035033