baidu88vip/pruning_yolov3

YOLOv3-training-prune

环境

Python3.6, Pytorch 1.0及以上， numpy>1.16， tensorboard=1.13以上

YOLOv3 的训练参考[博客]（https://blog.csdn.net/qq_34795071/article/details/90769094）， 代码基于的 PyTorch-YOLOv3

正常训练（Baseline）

python train.py --data data/VHR.data --cfg cfg/yolov3.cfg --weights/yolov3.weights --epochs 100 --batch-size 32 #后面的epochs自行更改 直接加载weights可以更好的收敛

剪枝算法介绍

本代码基于论文 Learning Efficient Convolutional Networks Through Network Slimming (ICCV 2017) 进行改进实现的 channel pruning算法，类似的代码实现还有这个 yolov3-network-slimming。原始论文中的算法是针对分类模型的，基于 BN 层的 gamma 系数进行剪枝的。

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

进行稀疏化训练

python train.py --cfg cfg/yolov3.cfg --data data/VHR.data --weights weights/last.pt --epochs 100 --batch-size 32 -sr --s 0.0001  #scale参数默认0.001，在数据分布广类别多的或者稀疏时掉点厉害的适当调小s

训练过程中模型可视化

tensorboard --logdir=runs 

模型剪枝

python prune.py --cfg cfg/yolov3.cfg --data data/my_data.data --weights weights/last.pt --percent 0.5

模型进行微调

python train.py --cfg cfg/prune_0.5_yolov3_cfg.cfg --data data/VHR.data --weights weights/prune_0.5_last.weights --epochs 100 --batch-size 32

convert cfg/pytorch model to darknet weights

python  -c "from models import *; convert('cfg/yolov3.cfg', 'weights/yolov3.pt')"
Success: converted 'weights/yolov3.pt' to 'converted.weights'

参考

https://github.com/Lam1360/YOLOv3-model-pruning

https://github.com/ultralytics/yolov3