YOLOX剪枝

YOLOX剪枝(官网的)

本文主要是根据csdn 爱吃肉的鹏，他的pytorch-yolox剪枝代码所改，非常感谢大佬的帮助，正好是我所需要的，但是我发现B站大佬这个代码训练出来的精度会比官网低10%的map 0.5，在我的Dota-1.5数据集，有一个类总是AP=0，我感觉应该是B站大佬的代码写的不完整，导致推理的精度低，所以我直接用的YOLOX官网，与大佬的版本相结合。又加了一个YOLOX的整体剪枝。

环境

https://github.com/Megvii-BaseDetection/YOLOX.git
这里把YOLOX的代码下载到本地，并根据其安装好环境，这里就不再赘述了。

安装包

pip install torch_pruning==0.2.7

本文实现的功能

1、支持单个卷积剪枝

2、支持网络层剪枝

3、支持模型的整体剪枝

4、剪枝后微调训练

5、修改了激活函数 silu->mish(因为项目用)

数据集格式：采用VOC数据集格式

网络剪枝

参考论文:Pruning Filters for Efficient ConvNets
首先进行模型的正常训练

1、修改代码

网络剪枝

参考论文:Pruning Filters for Efficient ConvNets
在剪枝之前需要通过tools/prunmodel.py save_whole_model(weights_path,num_classes)函数将模型的权重和结构都保存下来
weights_path:权重路径
num_classes:自己数据集的类别数

支持对某个卷积的剪枝 调用Conv_pruning(whole_model_weights):

pruning_idxs=strategy(v,amout=0.4) #0.4是剪枝率，根据需要自己修改，数越大剪枝的越多
对于单独一个卷积的剪枝，需要修改两个地方，这里的卷积层需要打印模型获得，不要自己盲目瞎猜:

支持网络层的剪枝:调用layer_pruning(whole_model_weights):

剪枝以后，会打印模型的参数量变化

支持模型整体剪枝!

python exps/network_slim/main.py
这里面的total_step=1表示只进行一轮剪枝
ch_sparsity=0.5表示剪枝率
我测试了yolox-s

剪枝后的微调训练

在tools/train.py中加入

然后在yolox/core/trainer.py中

这里的model_path换成自己剪枝后保存下来的pth文件就可以。

代码

https://github.com/lizexu123/YOLOX-pruning.git
点个关注把