Skip to content

bot66/yolov5_pruning

 
 

Repository files navigation

Yolov5 Pruning

Using "Learning Efficient Convolutional Networks Through Network Slimming (ICCV2017) channel pruning method prune yolov5s on MSCOCO 2017 Person dataset.

The Dataset

MSCOCO 2017 dataset. Reserve only no crowed person class and remove all other class label.

You can download MSCOCO 2017 dataset and use scripts/create_yolov5_dataset_from_coco.py to create the dataset.

The Model

Default yolov5s.

Results

The results are generated by test.py, Notice the speed is inconsistent with production level.

Network Sparsity Rate mAPval (without finetune)
0.5
mAPval(finetune)
0.5
mAPval(finetune)
0.5:0.95
Speed
RTX3070 (ms)
params
(M)
yolov5s 62.1 45.2 3.0 7.0
yolov5s 30% pruned 1e-5 10.5 62.4 44.8 2.7 3.5
yolov5s 50% pruned 1e-5 0 58.8 41.5 2.3 1.8
yolov5s 70% pruned 1e-5 0 46.6 30.5 1.8 0.7

Usage

Set pruning ratio as 0.7 for example.

Step1: Sparsity training from scrath.

python train.py --data dataset/data.yaml  --cfg models/yolov5s-person.yaml --epochs 150 --slimming --name yolov5s_coco-person_slimming --weights ''  --batch-size 32 

Step2: Prune the sparsed net.

python prune_net.py --weight /path/to/the/sparsed/weight --pruning_ratio 0.7

The pruned net weight and cfg file will save to "pruned_net" directory.

Step3: Finetune the pruned net.

python train.py --data dataset/data.yaml --hyp data/hyps/hyp.finetune.yaml --cfg pruned_net/prune_0.7_pruned_net.yaml --epochs 30  --name yolov5s_coco-person_slimming_finetune_0.7 --weights pruned_net/prune_0.7_pruned_net.pt  --batch-size 32 

Reference

@InProceedings{Liu_2017_ICCV,
    author = {Liu, Zhuang and Li, Jianguo and Shen, Zhiqiang and Huang, Gao and Yan, Shoumeng and Zhang, Changshui},
    title = {Learning Efficient Convolutional Networks Through Network Slimming},
    booktitle = {The IEEE International Conference on Computer Vision (ICCV)},
    month = {Oct},
    year = {2017}
}

foolwood/pytorch-slimming

Packages

No packages published

Languages

  • Python 98.1%
  • Shell 1.2%
  • Dockerfile 0.7%