pruning_yolov3's Introduction

YOLOv3-training-prune

环境

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

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

大家也可以看下这个：https://github.com/tanluren/yolov3-channel-and-layer-pruning

正常训练（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 参数的尝试或者需要进行迭代式剪枝等。

参数设置

-sr开启稀疏化，

--s指定稀疏因子大小，

--prune指定稀疏类型，

--prune 0为正常剪枝和规整剪枝的稀疏化

--prune 1为极限剪枝的稀疏化

进行稀疏化训练

baseline后生成pt 如果你不想加载yolov3.weights进行稀疏化训练把baseline最后的pt转化为weights 进行稀疏化

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

训练过程中模型可视化

tensorboard --logdir=runs

模型剪枝

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

shortcut_prune剪枝

python shortcut_prune.py  --cfg/yolov3.cfg --data data/VHR.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

pruning_yolov3's People

Contributors

Stargazers

Watchers

pruning_yolov3's Issues

YOLOV4_PRUNING

Hi ,

Is this repository compatible with yolov4 darknet weights (yolov4.weights) for prunning ?

torch.nn.modules.module.ModuleAttributeError: 'DistributedDataParallel' object has no attribute 'module_list'

我在执行正常训练的指令的时候（没有剪枝），报下面这个错误，请问您知道是什么原因吗？
plt.subplot(ns, ns, i + 1).imshow(imgs[i].transpose(1, 2, 0))
0%| | 0/297 [00:07<?, ?it/s]
Traceback (most recent call last):
File "train_pru.py", line 457, in
train() # train normally
File "train_pru.py", line 306, in train
BNOptimizer.updateBN(sr_flag, model.module_list, opt.s, prune_idx)
File "/home/share/lib/python3.7/site-packages/torch/nn/modules/module.py", line 772, in getattr
type(self).name, name))
torch.nn.modules.module.ModuleAttributeError: 'DistributedDataParallel' object has no attribute 'module_list'

zbyuan / pruning_yolov3 Goto Github PK