chendl's picture
Add application file
0b7b08a
|
raw
history blame
3.18 kB

YOLOX-ONNXRuntime in Python

This doc introduces how to convert your pytorch model into onnx, and how to run an onnxruntime demo to verify your convertion.

Step1: Install onnxruntime

run the following command to install onnxruntime:

pip install onnxruntime

Step2: Get ONNX models

Users might download our pre-generated ONNX models or convert their own models to ONNX.

Download ONNX models.

Model Parameters GFLOPs Test Size mAP Weights
YOLOX-Nano 0.91M 1.08 416x416 25.8 github
YOLOX-Tiny 5.06M 6.45 416x416 32.8 github
YOLOX-S 9.0M 26.8 640x640 40.5 github
YOLOX-M 25.3M 73.8 640x640 47.2 github
YOLOX-L 54.2M 155.6 640x640 50.1 github
YOLOX-Darknet53 63.72M 185.3 640x640 48.0 github
YOLOX-X 99.1M 281.9 640x640 51.5 github

Convert Your Model to ONNX

First, you should move to by:

cd <YOLOX_HOME>

Then, you can:

  1. Convert a standard YOLOX model by -n:
python3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth

Notes:

  • -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nano, yolox-tiny, yolov3]

  • -c: the model you have trained

  • -o: opset version, default 11. However, if you will further convert your onnx model to OpenVINO, please specify the opset version to 10.

  • --no-onnxsim: disable onnxsim

  • To customize an input shape for onnx model, modify the following code in tools/export.py:

    dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])
    
  1. Convert a standard YOLOX model by -f. When using -f, the above command is equivalent to:
python3 tools/export_onnx.py --output-name yolox_s.onnx -f exps/default/yolox_s.py -c yolox_s.pth
  1. To convert your customized model, please use -f:
python3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth

Step3: ONNXRuntime Demo

Step1.

cd <YOLOX_HOME>/demo/ONNXRuntime

Step2.

python3 onnx_inference.py -m <ONNX_MODEL_PATH> -i <IMAGE_PATH> -o <OUTPUT_DIR> -s 0.3 --input_shape 640,640

Notes:

  • -m: your converted onnx model
  • -i: input_image
  • -s: score threshold for visualization.
  • --input_shape: should be consistent with the shape you used for onnx convertion.