Skip to content

uyzhang/yolov5_prune

Repository files navigation

Introduction

Clean code version of YOLOv5(V6) pruning.

The original code comes from : https://github.com/midasklr/yolov5prune.

Steps:

  1. Basic training

    • In COCO Dataset
      python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 32 --device 0 --epochs 300 --name coco --optimizer AdamW --data data/coco.yaml
  2. Sparse training

    • In COCO Dataset
      python train.py --batch 32 --epochs 50 --weights weights/yolov5s.pt --data data/coco.yaml --cfg models/yolov5s.yaml --name coco_sparsity --optimizer AdamW --bn_sparsity --sparsity_rate 0.00005 --device 0
  3. Pruning

    • In COCO Dataset
      python prune.py --percent 0.5 --weights runs/train/coco_sparsity13/weights/last.pt --data data/coco.yaml --cfg models/yolov5s.yaml --imgsz 640
  4. Fine-tuning

    • In COCO Dataset
      python train.py --img 640 --batch 32 --epochs 100 --weights runs/val/exp1/pruned_model.pt  --data data/coco.yaml --cfg models/yolov5s.yaml --name coco_ft --device 0 --optimizer AdamW --ft_pruned_model --hyp hyp.finetune_prune.yaml

Experiments

  • Result of COCO Dataset

    exp_name model optim&epoch lr sparity [email protected] note prune threshold BN weight distribution Weight
    coco yolov5s adamw 100 0.01 - 0.5402 - - - -
    coco2 yolov5s adamw 300 0.01 - 0.5534 - - - last.pt
    coco_sparsity yolov5s adamw 50 0.0032 0.0001 0.4826 resume official SGD 0.54 -
    coco_sparsity2 yolov5s adamw 50 0.0032 0.00005 0.50354 resume official SGD 0.48 -
    coco_sparsity3 yolov5s adamw 50 0.0032 0.0005 0.39514 resume official SGD 0.576 -
    coco_sparsity4 yolov5s adamw 50 0.0032 0.001 0.34889 resume official SGD 0.576 -
    coco_sparsity5 yolov5s adamw 50 0.0032 0.00001 0.52948 resume official SGD 0.579 -
    coco_sparsity6 yolov5s adamw 50 0.01 0.0005 0.51202 resume coco 0.564 -
    coco_sparsity10 yolov5s adamw 50 0.01 0.001 0.49504 resume coco2 0.6 -
    coco_sparsity11 yolov5s adamw 50 0.01 0.0005 0.52609 resume coco2 0.6 -
    coco_sparsity13 yolov5s adamw 100 0.01 0.0005 0.533 resume coco2 0.55 last.pt
    coco_sparsity14 yolov5s adamw 50 0.01 0.0007 0.515 resume coco2 0.61 -
    coco_sparsity15 yolov5s adamw 100 0.01 0.001 0.501 resume coco2 0.54 -
  • The model of pruning coco_sparsity13

    coco_sparsity13 [email protected] Params/FLOPs
    origin 0.537 7.2M/16.5G
    after 10% prune 0.5327 6.2M/15.6G
    after 20% prune 0.5327 5.4M/14.7G
    after 30% prune 0.5324 4.4M/13.8G
    after 33% prune 0.5281 4.2M/13.6G
    after 34% prune 0.5243 4.18M/13.5G
    after 34.5% prune 0.5203 4.14M/13.5G
    after 35% prune 0.2548 4.1M/13.4G
    after 38% prune 0.2018 3.88M/13.0G
    after 40% prune 0.1622 3.7M/12.7G
    after 42% prune 0.1194 3.6M/12.4G
    after 45% prune 0.0537 3.4M/12.0G
    after 50% prune 0.0032 3.1M/11.4G