Yolov5 中添加Network Slimming剪枝--稀疏训练部分

最新推荐文章于 2024-08-10 15:52:57 发布

m0_46093829

最新推荐文章于 2024-08-10 15:52:57 发布

阅读量4.2k

点赞数 9

本文链接：https://blog.csdn.net/m0_46093829/article/details/128157589

版权

本文介绍了在Yolov5中应用Network Slimming剪枝方法进行模型压缩的过程，重点讲解了稀疏训练的核心思想——通过L1正则化约束BN层系数，以剪掉贡献小的通道。详细分析了添加L1正则化的代码实现，并解释了为何只对特定BN层进行操作的原因。同时，提供了查看稀疏训练效果的方法，展示了随着训练进行，权重分布的变化情况。

摘要由CSDN通过智能技术生成

前言：Network Slimming剪枝过程让如下

1. 稀疏化

2. 剪枝

3. 反复迭代这个过程

一、稀疏化：

通过Network Slimming 的核心思想是:添加L1正则来约束BN层系数，从而剪掉那些贡献比较小的通道channel

原理如下：BN层的计算是这样的：

上边介绍了，Network Slimming的核心思想是剪掉那些贡献比较小的通道channel，它的做法是从BN层下手。BN层的计算公式如下：

通过BN层的计算公式可以看出每个channe的Zout的大小和系数γ正相关，因此我们可以拿掉哪些γ-->0的channel，但是由于正则化，我们训练一个网络后，bn层的系数是正态分布的。这样的话，0附近的值则很少，那剪枝的作用就很小了。因此要先给BN层加上L1正则化进行一步稀疏训练（为什么要用L1正则化可以看该博客：l1正则与l2正则的特点是什么，各有什么优势？ - 知乎）。

为BN层加入L1正则化后，损失函数公式为：

上面第一项是正常训练的loss函数，第二项是约束对于L1正则化，g(s)=|s|，λ是正则系数，引入L1正则来控制γ, 要把稀疏表达加在γ 上, 得到每个特征的重要性 λ

- 每个通道的特征对应的权重是 γ 
- 稀疏表达也是对 γ 来说的, 所以正则化系数 λ 也是针对  γ, 而不是 W
-  稀疏化后, 做γ 值的筛选

因此在进行反向传播时候：𝐿′=∑𝑙′+𝜆∑𝑔′(𝛾)=∑𝑙′+𝜆∑|𝛾|′=∑𝑙′+𝜆∑𝑠𝑖𝑔𝑛(𝛾)

那如何把程序加到yolov5呢？

在yolov5 train.py的程序中找到反向传播部分程序：

1.1 稀疏训练核心代码

将scaler.scale(loss).backward()注释，并添加下方代码：

代码如下：

 # Backward
            # scaler.scale(loss).backward()
            loss.backward()
            # # ============================= sparsity training ========================== #
            srtmp = opt.sr*(1 - 0.9*epoch/epochs)  # opt.sr=0.0001 随着epoch增多，把srtmp减小
            if opt.st:  # '默认是true  train with L1 sparsity normalization  
                ignore_bn_list = []
                for k, m in model.named_modules():
                    # print('name: {}, module: {}'.format(k,m))
                    if isinstance(m, Bottleneck):
                        if m.add:
                            ignore_bn_list.append(k.rsplit(".", 2)[0] + ".cv1.bn")
                            ignore_bn_list.append(k + '.cv1.bn')