【模型剪枝2】不同剪枝方法实现对 yolov5n 剪枝测试及对比

目录

一、背景

二、剪枝

1. Network Slimming

1.0 代码准备

1.1 稀疏化训练

1.2 剪枝

1.3 微调

1.4 测试总结

2. Torch Pruning（TP）

2.1 MagnitudePruner

2.1.1 剪枝

2.1.2 retrain

2.1.3 测试总结

2.2 SlimmingPruner

2.2.1 定义重要性评估

2.2.2 定义剪枝器

2.2.3 初始化

2.2.4 稀疏化训练

2.2.5 剪枝

2.2.6 微调

2.2.7 测试总结

三、总结

参考资料

---

一、背景

        书接上回，本文记录一下上文中介绍的三种剪枝方法的使用，以剪枝 yolov5n6 为例，对比一下剪枝前后模型体积和准确率变化，以及不同剪枝方法之间的对比。

二、剪枝

        [1] 官方没有提供代码，放到后面使用 Torch Pruning 工具实现，[2] 的官方代码是对 VGG 和 ResNet 进行的剪枝，有需要的朋友可以参考，这里使用的是 midasklr 大佬提供的 yolov5 实现[5]，也是使用比较多的。

1. Network Slimming

        主要剪枝流程如下，作为比较，可以先训练一个原始 yolo 模型作为 baseline，之后再进行稀疏化训练。

剪枝流程图

原始 yolov5n6 模型精度测试结果

1.0 代码准备

        首先本文测试的模型是有 4 个检测尺度分支的 yolov5n6，所以需要先对源码进行修改已适配模型结构。

（1）在 models/yolo.py 的 Line538 后添加一行

from_to_map[named_m_base + ".m.3"] = fromlayer[f[3]]

（2）prune.py 从 Line432 开始的 pruned_yaml["backbone"] 和 pruned_yaml["head"] 改成 yolov5n6 的 backbone 和 head

（3）Line554 的 "model.24" 改成 “model.33”

1.1 稀疏化训练

        train_sparity.py 提供了稀疏化训练代码，这里的稀疏率  没有像官方一样使用定值，而是根据训练轮次增加逐渐减小，通常  越大稀疏化程度越高。

python train_sparity.py --st --sr 0.0005

训练完成后可以进入 http://localhost:6006/ 查看 BN 层权重分布直方图的变化情况

tensorboard --logdir=./runs/train

稀疏困难

稀疏过快

平稳变化

         如果设置过小，BN 层很难稀疏化，可剪枝空间很小；如果设置过大，BN 层会过于稀疏化，难以学习到有效权重，造成精度损失和过度剪枝。所以需要根据具体任务和训练结果不断调整 ，直到 BN 层的权重在训练过程中平稳变化至 0 点。

除了监控 BN 层，还可以通过 mAP 的变化情况判断模型的训练效果。

稀疏模型的精度测试

1.2 剪枝

稀疏训练完成后对得到的稀疏模型进行剪枝

python prune.py --percent 0.5 --weights /path/to/your/sparse_model

可以看到根据计算得到的  阈值对各层进行了剪枝。

        当剪枝之后剩余通道数为 1 时，剪枝过程可能会报错。上图中 model.8.cv3.bn 剪枝后剩余 1 个通道，然后就会报错。

这表明对该层来说  阈值偏大了，所以需要针对该层调整 ，确保剪枝后剩余通道数＞1，具体需要修改 utils/prune_utils.py 的 obtain_bn_mask()，每次剪完之后判断剩余通道数，如果余 1，则按1% 的步长降低 ，这样能保证剪枝通道尽可能多，循环剪枝判断，直至剩余通道数＞1。

此时再剪枝就没有问题了，可以看到 model.8.cv3.bn 层剩余通道数变2了。

剪枝 50% 后的精度测试

可以看到跟剪枝前的稀疏模型相比没有精度损失，说明还有剪枝空间，剪枝率可以再调大一些。

剪枝 70% 后的模型和精度测试：

可以看到准确率明显下降，需要微调来恢复精度。 

1.3 微调

python finetune.py --weights pruned_model.pt

微调 60 个 epochs 之后准确率恢复到可接受程度，相比剪枝前会有轻微下降。