Github项目推荐 | YOLOv3模型压缩

• 2020 年 8 月 10 日
• AI
• 目标检测

简介

项目仓库包括具有一些轻量级主干（ShuffleNetV2，GhostNet），一些计算机视觉注意机制（SE Block, CBAM Block, ECA Block），针对GhostNet的剪枝、量化和提炼的YOLOv3。

项目地址

//github.com/HaloTrouvaille/YOLO-Multi-Backbones-Attention

重要更新记录

2020.06.01

1）最佳的轻量级模型-HuaWei GhostNet被添加为YOLOv3骨干！ 它比ShuffleNetV2更好。 visdrone数据集的结果如下。

2）添加Dorefa量化方法以进行任意位量化！ visdrone数据集的结果如下。

3）删除了ShuffleNet和注意力机制。

2020.06.24

1）根据NetworkSlimming添加修剪。

2）修剪后添加提炼以获得更高的mAP。

3）为GhostNet添加Imagenet预训练模型。

注意：单GPU的使用效果更佳。

如果需要以前的注意力模型或有任何疑问，可以到项目中联系作者。

运行环境

• python 3.7

• pytorch >= 1.1.0

• opencv-python

数据集

• Oxfordhand datasets (1 class, including human’s hand)

//pan.baidu.com/s/1ZYKXMEvNef41MdG1NgWYiQ (extract code: 00rw)

• Visdrone Remote datasets (10 classes, including pedestrian, car, bus, etc)

//pan.baidu.com/s/1JzJ6APRym8K64taZgcDZfQ (extract code: xyil)

• Bdd100K datasets (10 classes, including motor, train, traffic light, etc)

//pan.baidu.com/s/1dBrKEdy92Mxqg-JiyrVjkQ (extract code: lm4p)

• Dior datasets (20 classes, including airplane, airport, bridge, etc)

//pan.baidu.com/s/1Fc-zJtHy-6iIewvsKWPDnA (extract code: k2js)

使用

1）下载数据集，将其放置在数据（data）目录中

2）使用以下命令训练模型（通过更改cfg文件来更改模型结构）

python3 train.py --data data/visdrone.data --batch-size 16 --cfg cfg/ghostnet-yolov3-visdrone.cfg --img-size 640

3）使用训练有素的模型检测对象（将图片或视频放置在示例目录中）

python3 detect.py --cfg cfg/ghostnet-yolov3-visdrone.cfg --weights weights/best.pt --data data/visdrone.data

4）结果

剪枝和量化

剪枝

首先，执行稀疏训练。

python3 train.py --data data/visdrone.data --batch-size 4 --cfg cfg/ghost-yolov3-visdrone.cfg --img-size 640 --epochs 300  --device 3 -sr --s 0.0001

然后在normal_prune.py中更改cfg和权重，再使用以下命令

python normal_prune.py

获取pruned.cfg和相应的权重文件后，可以通过以下命令微调修剪的模型

python3 train.py --data data/visdrone.data --batch-size 4 --cfg pruned.cfg --img-size 640 --epochs 300  --device 3 --weights weights/xxx.weighs

量化

如果要量化某些卷积层，可以在cfg文件中将[convolutional]更改为[quan_convolutional]。 然后使用以下命令

python3 train.py --data data/visdrone.data --batch-size 16 --cfg cfg/ghostnet-yolov3-visdrone.cfg --img-size 640

更改YOLOv3主干的实验结果

ShuffleNetV2+双标度检测（YOLO检测器）

使用Oxfordhand数据集

使用Visdrone数据集（不完整的训练）

注意机制的实验结果

基于YOLOv3-tiny

SE Block paper : //arxiv.org/abs/1709.01507

CBAM Block paper : //arxiv.org/abs/1807.06521

ECA Block paper : //arxiv.org/abs/1910.03151

TODO

√ ShuffleNetV2主干

√ 华为GhostNet骨干网

√ ImageNet预训练

√ COCO数据集培训

□ 其他检测策略

□ 其他剪枝策略