基于Pytorch框架下的YOLOV3剪枝算法实现

代码：https://github.com/Lam1360/YOLOv3-model-pruning
算法：https://arxiv.org/abs/1708.06519 ，代码基于论文 Learning Efficient Convolutional Networks Through Network Slimming (ICCV 2017) 进行改进实现的 channel pruning算法，类似的代码实现还有这个 yolov3-network-slimming。原始论文中的算法是针对分类模型的，基于 BN 层的 gamma 系数进行剪枝的。

1、环境：

python3.6，Pytorch 1.0及以上，YOLOv3 的实现参考了 eriklindernoren 的 PyTorch-YOLOv3 ，因此代码的依赖环境也可以参考其 repo；

用 YOLOv3 模型在一个开源的人手检测数据集 oxford hand 上做人手检测，并在此基础上做模型剪枝。对于该数据集，对 YOLOv3 进行 channel pruning 之后，模型的参数量、模型大小减少 80% ，FLOPs 降低 70%，前向推断的速度可以达到原来的 200%，同时可以保持 mAP 基本不变。

2、数据集：

①下载数据集，得到压缩文件
②将压缩文件解压到 data 目录，得到 hand_dataset 文件夹
③在 data 目录下执行 converter.py，生成 images、labels 文件夹和 train.txt、valid.txt 文件。训练集中一共有 4807 张图 片，测试集中一共有 821 张图片

3、权重文件：

把需要的权重文件放入weights文件夹，如yolov3.weights等

4、正常训练：

命令：

$ python train.py --model_def config/yolov3-hand.cfg 

5、剪枝算法的大概步骤以下只是算法的大概步骤，具体实现过程中还要做 s 参数的尝试或者需要进行迭代式剪枝等。

① 进行稀疏化训练

$ python train.py --model_def config/yolov3-hand.cfg -sr --s 0.01

②基于 test_prune.py 文件进行剪枝，得到剪枝后的模型python test_prune.py


剪枝后在checkpoints文件生成了剪枝的权重文件，在config文件夹生成了新的剪枝配置文件。


③剪枝后重新训练进行微调：

$ python train.py --model_def config/prune_0.85_yolov3-hand.cfg -pre checkpoints/prune_0.85_yolov3_ckpt_99_09240839.pth

④剪枝后模型的测试:

$ python test.py --model_def config/prune_0.85_yolov3-hand.cfg -pre checkpoints/prune_0.85_yolov3_ckpt_99_09240839.pth --data_config config/oxfordhand.data --class_path data/oxforhand.names --conf_thres 0.01