- 2021.12.15
- change backbone to Ghostnet
- Finish EagleEye pruning YOLOv5 series
- 2021.12.27
- change backbone to shufflenetv2
- change backbone to efficientnetv2
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
Models | [email protected] | [email protected]:.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