yolov5-6.0的稀疏化训练、结构化剪枝、微调

本文针对yolov5n的模型进行结构化剪枝、微调，在文章最后会放上程序链接。
剪枝方法：基于BN层系数gamma剪枝。
前言：yolov5n模型是非常轻量级的网络模型，但针对嵌入式端依然比较大；此时，可以通过添加L1正则来约束BN层系数，使得系数稀疏化，通过稀疏训练后，裁剪掉稀疏很小的层，对应激活也很小，所以对后面的影响非常小，反复迭代这个过程，可以获得很compact的模型。
原理：
在一个卷积-BN-激活模块中，BN层可以实现通道的缩放。如下：

BN层的具体操作有两部分：

在归一化后会进行线性变换，那么当系数gamma很小时候，对应的激活（Zout）会相应很小。这些响应很小的输出可以裁剪掉，这样就实现了bn层的通道剪枝。
通过在loss函数中添加gamma的L1正则约束，可以实现gamma的稀疏化。

上面损失函数L右边第一项是原始的损失函数，第二项是约束，其中g(s) = |s|，λ是正则系数，根据数据集调整。
实际训练的时候，就是在优化L最小，依据梯度下降算法：

′=∑′+∑′()=∑′+∑||′=∑′+∑∗()

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

            # Backward
            loss.backward()
            # scaler.scale(loss).backward()
            # # ============================= sparsity training ========================== #
            srtmp = opt.sr*(1 - 0.9*epoch/epochs)
            if opt.st:
                ignore_bn_list = []
                for k, m in model.named_modules():
                    if isinstance(m, Bottleneck):
                        if m.add:
                            ignore_bn_list.append(k.rsplit(".", 2)[0] + ".cv1.bn")
                            ignore_bn_list.append(k + '.cv1.bn')
                            ignore_bn_list.append(k + '.cv2.bn')
                    if isinstance(m, nn.BatchNorm2d) and (k not in ignore_bn_list):
                        m.weight.grad.data.add_(srtmp * torch.sign(m.weight.data))  # L1
                        m.bias.grad.data.add_(opt.sr*10 * torch.sign(m.bias.data))  # L1
            # # ============================= sparsity training ========================== #

            optimizer.step()
                # scaler.step(optimizer)  # optimizer.step
                # scaler.update()
            optimizer.zero_grad()

这里并未对所有BN层gamma进行约束，这里对C3结构中的Bottleneck结构中有shortcut的层不进行剪枝，主要是为了保持tensor维度可以加：

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

具体步骤：
1、首先，下载yolov5-6.0（其他版本应该也支持，目前只尝试了6.0/6.1）的源码，下载yolov5n.pt文件，配置训练自己数据集，得到一个模型。
2、然后，进行稀疏化训练；
train_sparity.py
在train.py中添加如下代码：

            # Backward
            # scaler.scale(loss).backward()
            loss.backward()
            # # ============================= sparsity training ========================== #
            srtmp = opt.sr*(1 - 0.9*epoch/epochs)
            if opt.st:
                ignore_bn_list = []
                for k, m in model.named_modules():
                    if isinstance(m, Bottleneck):
                        if m.add:
                            ignore_bn_list.append(k.rsplit(".", 2)[0] + ".cv1.bn")
                            ignore_bn_list.append(k + '.cv1.bn')
                            ignore_bn_list.append(k + '.cv2.bn')
                    if isinstance(m, nn.BatchNorm2d) and (k not in ignore_bn_list):
                        m.weight.grad.data.add_(srtmp * torch.sign(m.weight.data))  # L1
                        m.bias.grad.data.add_(opt.sr*10 * torch.sign(m.bias.data))  # L1

设置好如下的参数，加载刚才训练的模型，即可稀疏化训练：
此时训练容易报错：
ModuleNotFoundError: No module named ‘utils.loggers.wandb’
将下面47行注释掉即可。

稀疏训练前后map值变化如下所示：

稀疏训练前的训练：
稀疏训练50轮：

效果有所增加。
3、训练完成后进行剪枝
prune.py
此时容易出现：TypeError: run() got an unexpected keyword argument 'cfg’错误，
解决方法：在prune.py的 def run 和 def run_prune 中添加模型文件和参数信息即可，切记不可只在 def parse_opt 中添加。
4、微调训练测试
在finetune_prune.py中设置模型和data路径后，运行：
python finetune_pruned.py --weights pruned_model.pt --adam --epochs 50

map0.5值为0.865，相比正常训练时的0.871，精度下降了0.006，但模型大小由3.57兆减小为2.64兆，大小下降了26%，精度下降可忽略不计。
剪枝模型测试：修改detectpruned.py中模型和img路径，运行即可。

程序网盘链接：
链接：https://pan.baidu.com/s/1bGUonuaUPdeoPY-6IXOGPQ?pwd=1234
提取码：1234