- 2021.12.15
- change backbone to Ghostnet
- Finish EagleEye pruning YOLOv5 series
- 2021.12.27
- change backbone to shufflenetv2
- change backbone to efficientnetv2
- 2022.05.23
- Finish network slimming pruning YOLOv5 series
- Finish network SFP pruning YOLOv5 series
- Finish network FPGM pruning YOLOv5 series
pip install -r requirements.txt
python train.py --data data/VisDrone.yaml --imgsz 640 --weights '' --cfg models/yolo_ghostnet.yaml --nosave --device 0,1,2,3 --sync-bn
You can change depth_multiple and width_multiple to choose different yolov5 verson
- Normal Training
python train.py --data data/VisDrone.yaml --imgsz 640 --weights yolov5s.pt --cfg models/yolov5s_pruning.yaml --nosave --device 0,1,2,3 --sync-bn
- Search for Optimal Pruning Network
python eagleeye.py --data data/VisDrone.yaml --weight the_first_step_trained_model --cfg models/yolov5_pruning.yaml --path models/yolov5s_pruned.yaml --pruned_weights pruned_weight.pt
- Fine-tuning
python train.py --data data/VisDrone.yaml --imgsz 640 --weights pruned_weight.pt --cfg models/yolov5s_pruned.yaml --device 0,1,2,3 --nosave --sync-bn
cd yolov5_fpgm_slimming_sfp
- soft mask
python train.py --data data/VisDrone.yaml --imgsz 640 --weights yolov5s.pt --cfg models/yolov5s_pruning.
yaml --device 0,1,2,3 --sfp/fpgm --sfp_ratio/fpgm_ratio 0.5 --path models/yolov5s_pruned.yaml
- Fine_tuning
python train.py --data data/VisDrone.yaml --imgsz 640 --weights pruned_weights.pt --cfg models/yolov5s_fpgm/sfp_pruned.yaml --device 0,1,2,3
cd yolov5_fpgm_slimming_sfp
- BatchNorm Layer \gamma
python train.py --data data/VisDrone.yaml --imgsz 640 --weights yolov5s.pt --cfg models/yolov5s_pruning.yaml --device 0,1,2,3 -sr
- BatchNorm Layer pruning
python prune_slimming.py --weights the_first_step_trained_model --data data/VisDrone.yaml --device 0
- Fine_tuning
python train.py --data data/VisDrone.yaml --imgsz 640 --weights the_second_step_get_model --cfg models/yolov5s_slimm_pruned.yaml --device 0,1,2,3
| Models | mAP@.5 | mAP@.5:.95 | GFLOPS | Parameters(M) |
|---|---|---|---|---|
| yolov5s | 35.1 | 19.4 | 15.9 | 14.4 |
| yolov5l_Ghostnet | 33.1 | 18.2 | 42.7 | 49.4 |
| yolov5l_efficientnetv2 | 23.3 | 11.4 | 35.3 | 42.8 |
| yolov5L_shufflenetv2 | 29.0 | 15.2 | 38.0 | 40.2 |
| yolov5s_eagleeye | 30.0 | 15.5 | 8.6 | 8.0 |
- backbone: ShuffleNetV2
- backbone: EfficientNetV2
- backbone: SwinTrans
- Prune: Other Algorithms
- Quantization
- Knowledge Distillation