/YOLOX-Slim

YOLOX is a high-performance anchor-free YOLO, exceeding yolov3~v5 with MegEngine, ONNX, TensorRT, ncnn, and OpenVINO supported. Documentation: https://yolox.readthedocs.io/

Primary LanguagePythonApache License 2.0Apache-2.0

YOLOX-Slim

Implement Learning Efficient Convolutional Networks Through Network Slimming on YOLOX

  • Sparse-training -> pruning -> fine-tuning
  • Automatic applying network slimming on model without modifying the model code(backbone/fpn/head)
  • Export pruned model as onnx model
  • Find good training recipe on COCO:
    • sparsity train epoch/lr
    • hyper parameter of sparsity train: s
    • scheduler of s: In the original paper the author used a fixed s parameter, perhaps using a scheduler could get better results
    • fine-tuning epoch/lr/ratio
  • Make train/fine-tuning with fp16 work

Experiment

Using YOLOX-s model for sparse training, pruning and then fine-tuning, I expected to be able to train a slimming model with better mAP and smaller parameters than YOLOX-Tiny. However, due to limited GPU resources, both models were not fully trained (120 and 25 rounds), and the training hyper parameters and strategies are yet to be explored.

Model size mAPval
0.5:0.95		 Params
(M)		 FLOPs
(G)		 weights notes
YOLOX-s(sparse-training) 640 37.9 9.0 26.8 github max_epoch 120, linear warm up to s=0.0001
YOLOX-s(slimming model) 640 16.5 1.58 6.43 github github-onnx network_slim_ratio=0.65, max_epoch=25
YOLOX-s(sparse-training) 640 33.8 9.0 26.8 github max_epoch 80, linear warm up 10 epoch s=0.0002
YOLOX-s(slimming model) 640 26.74 github network_slim_ratio=0.6, max_epoch=80
YOLOX-s(Official) 640 40.5 9.0 26.8 github
YOLOX-Tiny(Official) 416 32.8 5.06 6.45 github

Installation

Step1. Setup YOLOX

Step2. Install network slimming library:

pip3 install git+https://github.com/Sanster/pytorch-network-slimming.git@0.2.0

Quick Start

Generate pruning schema:

python3 tools/gen_pruning_schema.py --save-path ./exps/network_slim/yolox_s_schema.json --name yolox-s

Sparse-training, network_slim_sparsity_train_s is a hyper parameter that needs to be adjusted according to your data

python3 tools/train.py -d 4 -b 64 \
-f exps/network_slim/yolox_s_slim_train.py \
-expn yolox_s_slim_sparsity_train \
network_slim_sparsity_train_s 0.0001

Apply network slimming and fine-tuning pruned model

python3 tools/train.py -d 4 -b 64 \
-f exps/network_slim/yolox_s_slim.py \
-c ./YOLOX_outputs/yolox_s_slim_sparsity_train/latest_ckpt.pth \
-expn yolox_s_slim_fine_tuning \
network_slim_ratio 0.65

Use slimming model nun demo.py

python3 tools/demo.py image \
-f exps/network_slim/yolox_s_slim.py \
-c ./YOLOX_outputs/yolox_s_slim_fine_tuning/latest_ckpt.pth \
--path assets/dog.jpg  --save_result

Use slimming model run eval.py

python3 tools/eval.py -d 4 -b 64 \
-f exps/network_slim/yolox_s_slim.py \
-c ./YOLOX_outputs/yolox_s_slim_fine_tuning/latest_ckpt.pth \
--conf 0.001

Deployment

Export pruned model as onnx model

python3 tools/export_onnx.py \
-f exps/network_slim/yolox_s_slim.py \
-c ./YOLOX_outputs/yolox_s_slim_fine_tuning/latest_ckpt.pth \
--output-name ./YOLOX_outputs/yolox_s_slim_fine_tuning/latest_ckpt.onnx

Run onnx_inference.py

python3 demo/ONNXRuntime/onnx_inference.py \
-m ./YOLOX_outputs/yolox_s_slim_fine_tuning/latest_ckpt.onnx \
-i assets/dog.jpg -o YOLOX_outputs