yolov3通道剪枝层剪枝知识蒸馏：无人机数据集案例

项目地址

https://github.com/tanluren/yolov3-channel-and-layer-pruning

前言

这是GitHub工程yolov3-channel-and-layer-pruning的一个案例，不少朋友已经通过这个工程完成了针对自己数据集的模型压缩，这里我们将通过无人机数据集visdrone再演示一下模型剪枝和蒸馏的过程，并且对比slimyolov3的实验结果，我们的精度和压缩度远超过它。项目支持yolov3和yolov3-spp，在基础训练，稀疏训练，剪枝，微调四个环节提供了完整且多样的策略，详细介绍可以看GitHub readme。讨论可以发issue或者加群734912150。

visdrone案例

首先下载数据集，标签格式转换参考了这里，编辑visdrone.data和visdrone.names文件，修改yolov3-spp.cfg的filters和classes得到visdrone-spp.cfg，官网下载yolov3-spp.weights到weights文件夹。因为时间和资源有限，我们使用了apex混合精度来加速训练，混合精度结果可能比全精度差丢丢，项目里有apex的介绍。

基础训练

基础训练的目的是得到一个高精度的大模型，它的精度和参数选择都对后续的稀疏训练和微调训练有参考意义，同时也可以用于知识蒸馏。不建议跳过基础训练直接稀疏。我们修改了load_darknet_weights的规则，读取yolov3.weights和yolov3-spp.weights时会跳过三个输出卷积，用其余部分来初始化模型的backbone和head，这有助于模型快速收敛（想使用原始权重推理的就要改名了）。在尝试了几个姿势后我们选择了修改学习率为默认十分之一来跑这次实验。

python train.py --cfg cfg/visdrone-spp.cfg --data data/visdrone.data --weights weights/yolov3-spp.weights --epochs 120 --batch-size 20

这次的推理分辨率是416，模型大小239MB，训练情况如下，得到的baseline mAP为0.173，visdrone是一个非常有挑战性的数据集。

看看现在的bn情况，因为未压缩，总体上近似正态分布。

稀疏训练

最关键一步，参考GitHub里oxfordhand案例来设置稀疏因子s，但因为我们学习率降为了十分之一，s应该增大十倍才能达到相当的压缩速度，hand案例大概50epoch达到较大压缩，为使稀疏后精度保持较好，s只增大了五倍，为0.005，预计100epoch达到较大压缩，总epochs定为300，确保后面有足够时间调整。

先转换一下baseline的权重为darknet的weights格式，得到converted.weights，以去掉其中的epoch信息。

python -c "from models import *; convert('cfg/visdrone-spp.cfg', 'weights/last.pt')"

然后开启稀疏训练。稀疏策略为策略一，恒定s，也就是添加的额外梯度大小保持不变。

python train.py --cfg cfg/visdrone-spp.cfg --data data/visdrone.data --weights converted.weights --batch-size 20 --epochs 300 -sr --s 0.005 --prune 1

稀疏的结果我们先看看bn图，压缩速度和我们预期差不多，大概在100epoch完成了较大压缩，后续都是细微调整。（因为训练被中断过，图分成了两部分）

再看看mAP和loss的变化，可以看到loss在bn大幅压缩的阶段也在不断上升，而mAP在不断下降，之后进入一段调整，而在210epoch左右loss快速下降，mAP也较大回升，注意此时bn还是处于压缩状态的，这是我们设置了epochs0.7阶段进行学习率衰减，模型精度得到修复。最终稀疏后mAP0.143,，这个结果不是很好，相比baseline掉了三个点，对后面剪枝和微调可能不利，但这里正好演示下微调和蒸馏的作用。

剪枝

稀疏好了，先来试试剪通道，通道策略三一般效果最好，直接用它了。随便剪一个

python slim_prune.py --cfg cfg/visdrone-spp.cfg --data data/visdrone.data --weights weights/last.pt --global_percent 0.9 --layer_keep 0.01

全局剪0.9，层保留0.01，剪后模型8.93MB，但是掉点厉害，不要。再降低比例试试

python slim_prune.py --cfg cfg/visdrone-spp.cfg --data data/visdrone.data --weights weights/last.pt --global_percent 0.85 --layer_keep 0.01

全局剪0.85，层保留0.01，剪后模型12.8MB，精度相比稀疏完掉点很少，而参数压缩了95%。这个还不错，拿来试试剪层

python layer_prune.py --cfg cfg/prune_0.85_keep_0.01_visdrone-spp.cfg --data data/visdrone.data --weights weights/prune_0.85_keep_0.01_last.weights --shortcuts 16

剪了16个shortcut，相当于剪了48个层，剪后模型11.3MB，精度掉了两点多，相比还在接受范围内，再看看能不能掉点少点

这次剪12个shortcut，模型12.2MB，精度相比之前掉了不到一个点。看多几个



可以这样去尝试不同的比例，同时考虑速度与精度的trade off，找到合适的剪法。大家可能会疑惑马赛克的是什么，其实那是模型推理速度，这个因机器而异，剪通道可能在某些机器没有加速效果，但剪层基本都能加速。剪枝后的模型在参数量，计算量，内存占用，推理速度方面都有很大优势，后面会分享案例中的权重和cfg给大家测试。

微调

前面剪枝的模型出现了不同程度的掉点，这时我们要进行微调恢复模型精度。有些朋友剪枝后模型精度反而还高了，可能会忽略微调这步，但剪枝毕竟令模型发生了突变，计算会产生偏差，可能会有一些意料之外的检测结果，微调还是有必要的。微调可以在U版这里跑，也可以拿cfg和权重到AB版darknet跑，效果也很好。
我们在微调这步增加了知识蒸馏策略，因为蒸馏在大小模型结构相似的情况下效果更明显，特别适用我们这里剪枝后的微调。
先来微调一个

python train.py --cfg cfg/prune_12_shortcut_prune_0.85_keep_0.01_visdrone-spp.cfg --data data/visdrone.data --weights weights/last.pt --batch-size 20 --epochs 100 --t_cfg cfg/visdrone-spp.cfg --t_weights weights/9visdrone-spp/last.pt

这个剪了12个shortcut结构，大小12.2MB，微调后精度由0.132恢复到了0.167。再看另一个

python train.py --cfg cfg/prune_16_shortcut_prune_0.85_keep_0.01_visdrone-spp.cfg --data data/visdrone.data --weights weights/last.pt --batch-size 20 --epochs 100 --t_cfg cfg/visdrone-spp.cfg --t_weights weights/9visdrone-spp/last.pt

这个权重剪了48个层，大小11.3M，微调后精度由0.118恢复到了0.163。更极端一点情况呢

python train.py --cfg cfg/prune_20_shortcut_prune_0.85_keep_0.01_visdrone-spp.cfg --data data/visdrone.data --weights weights/last.pt --batch-size 20 --epochs 100 --t_cfg cfg/visdrone-spp.cfg --t_weights weights/9visdrone-spp/last.pt

这个剪了20个shortcut的也回到了0.152。一般来说，想精度保持好点，可以剪少点；想速度快点，可以加大比例，特别是剪层。

slimyolov3是在visdrone2018跑的，我们是在2019跑的，baseline对比不太公平，但从剪枝的效果看，我们的实验比它好很多，参数压缩95%以上，精度保持也更好，而且还剪了层，速度绝对杠。有空会整理个对比列表。这里分享了这个实验相关的cfg和权重，有兴趣的朋友可以下载来测试。

更多

陆续会补充更多实验记录。

这是同样模型把推理分辨率设置为608后重新跑的结果，baseline0.259，稀疏后0.236，剪了一个14.9M的权重，微调后0.262甚至稍微超过了baseline。

待续。。。