Home

Awesome

YOLOv3-complete-pruning

本项目以ultralytics/yolov3为YOLOv3的Pytorch实现,并在YOLOv3-model-pruning剪枝的基础上,推出了4个YOLO-v3剪枝版本。(在此致谢两位)

此外,最近还更新了YOLO的1bit、4bit、8bit、16bit量化。

剪枝方式<center>优点</center><center>缺点</center>
正常剪枝不对shortcut剪枝,拥有可观且稳定的压缩率,无需微调。压缩率达不到极致。
极限剪枝极高的压缩率。需要微调。
规整剪枝专为硬件部署设计,剪枝后filter个数均为8的倍数,无需微调。为规整牺牲了部分压缩率。
Tiny剪枝稳定的压缩率。由于Tiny本来已很小,压缩率中规中矩。

项目特点

1.采用的YOLO-v3实现较为准确,mAP相对较高。

模型320416608
YOLOv351.8 (51.5)55.4 (55.3)58.2 (57.9)
YOLOv3-tiny29.032.9 (33.1)35.5

2.提供对YOLOv3及Tiny的多种剪枝版本、量化版本,以适应不同的需求。

3.剪枝后保存为.weights格式,可在任何框架下继续训练、推理,或以图像视频展示。

<img src="https://pjreddie.com/media/image/Screen_Shot_2018-03-24_at_10.48.42_PM.png" width="500">

4.目前支持情况

<center>剪枝方式</center><center>单卡</center><center>多卡</center>
<center>正常训练</center><center></center><center></center>
<center>稀疏化</center><center></center><center></center>
<center>正常剪枝</center><center></center><center></center>
<center>规整剪枝</center><center></center><center></center>
<center>极限剪枝(shortcut)</center><center></center><center></center>
<center>Tiny剪枝</center><center></center><center></center>
<center>二值量化</center><center>8bit量化</center><center>16bit量化</center><center>混合量化</center><center>任意bit量化</center>
<center></center><center></center><center></center><center></center><center></center>

最新进展

环境搭建

1.由于采用ultralytics/yolov3的YOLO实现,环境搭建见ultralytics/yolov3。这里重复介绍一下:

可直接pip3 install -U -r requirements.txt搭建环境,或根据该.txt文件使用conda搭建。

数据获取

依然采用oxford hand数据集

<img src="https://www.robots.ox.ac.uk/~vgg/data/hands/hand_dataset_files/Picture_1.png" width="500">

1.下载数据集,并解压至/data目录下,得到hand_dataset文件夹。

2.执行python converter.py ,生成 images、labels 文件夹和 train.txt、valid.txt 文件。

3.获取YOLO预训练权重,/weights文件夹下执行bash download_yolov3_weights.sh,或自行下载。

4.至此,数据部分完成。

剪枝篇

1.正常训练

python3 train.py --data data/oxfordhand.data --batch-size 32 --accumulate 1 --weights weights/yolov3.weights --cfg cfg/yolov3-hand.cfg

2.稀疏化训练

-sr开启稀疏化,--s指定稀疏因子大小,--prune指定稀疏类型。

其中:

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

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

--prune 2为Tiny剪枝的稀疏化

python3 train.py --data data/oxfordhand.data --batch-size 32 --accumulate 1 --weights weights/yolov3.weights --cfg cfg/yolov3-hand.cfg -sr --s 0.001 --prune 0 

3.模型剪枝

python3 normal_prune.py
python3 regular_prune.py
python3 shortcut_prune.py
python3 prune_tiny_yolo.py

需要注意的是,这里需要在.py文件内,将opt内的cfg和weights变量指向第2步稀疏化后生成的cfg文件和weights文件。 此外,可通过增大代码中percent的值来获得更大的压缩率。(若稀疏化不到位,且percent值过大,程序会报错。)

量化篇

1 指定需要量化的层

打开任意一个可用的配置文件,例如yolov3-hand.cfg。将需要量化的层,从原来的convolutional替换为quantize_convolutional。

2 指定量化方式

通过修改models.py中的W_bit和A_bit,指定权重的量化方式,激活的量化方式。(目前默认为16bit量化)

3 量化训练

python3 train.py --data data/oxfordhand.data --batch-size 32 --accumulate 1 --weights weights/yolov3.weights --cfg cfg/yolov3-quantize-hand.cfg

与正常训练相同,只是cfg要指向修改过的cfg文件。

推理展示

这里,我们不仅可以使用原始的YOLOV3用来推理展示,还可使用我们剪枝后的模型来推理展示。(修改cfg,weights的指向即可)

<img src="https://user-images.githubusercontent.com/26833433/64067835-51d5b500-cc2f-11e9-982e-843f7f9a6ea2.jpg" width="500">
python3 detect.py --source ...

例如:

python3 detect.py --cfg cfg/prune_0.8_yolov3-hand.cfg --weights weights/yolov3_hand_pruning_percent0.8.weights --data data/oxfordhand.data --source test.jpg

剪枝效果

以下数据除极限剪枝外,均未微调。

YOLO-v3剪枝

模型参数量模型体积压缩率耗时mAP
Baseline(416)61.5M246.4MB0%11.7ms0.7924
正常剪枝10.9M43.9MB82.2%5.92ms0.7712
规整剪枝15.31M61.4MB75.1%6.01ms0.7832
极限剪枝7.13M28.6MB88.4%5.90ms0.7382

YOLO-v3-Tiny剪枝

模型参数量模型体积压缩率耗时mAP
Baseline(416)8.7M33.1MB0%2.2ms0.6378
Tiny剪枝4.4M16.8MB40.1%2.0ms0.6132

量化效果(数据更新中)

以下量化数据中,权重激活均被量化。

YOLO-v3量化

模型mAP
Baseline(416)0.8246
第一层float32+中间层8bit+最后一层float320.8174
全16bit量化0.8132
全8bit量化0.8024

核心思想

剪枝方法来源于论文Learning Efficient Convolutional Networks through Network Slimming,剪枝无需微调方法来源于Rethinking the Smaller-Norm-Less-Informative Assumption in Channel Pruning of Convolution Layers

量化方法来源于论文DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients

此外,具体实现时,在论文作者的基础上做了改进。

互动

1.如何获得较高的压缩率?

提高压缩率的关键在于稀疏化训练,可以加大--s的值并迭代训练多次等手段。

2.我的压缩率比表格中更高!

以上数据仅仅是测试了不到20次的结果,如果有同学的压缩率更高,欢迎在评论区分享!

3.程序报错怎么办?

YOLOv3报错

由于采用了ultralytics/yolov3为YOLOv3的Pytorch实现,因此这类错误可跳转至此链接询问。

剪枝量化错误

一定要在本评论区留言,我会尽快修正!