YOLOX-TensorRT in Python¶

This tutorial includes a Python demo for TensorRT.

Install TensorRT Toolkit¶

Please follow the TensorRT Installation Guide and torch2trt gitrepo to install TensorRT and torch2trt.

YOLOX models can be easily conveted to TensorRT models using torch2trt

If you want to convert our model, use the flag -n to specify a model name:

python tools/trt.py -n <YOLOX_MODEL_NAME> -c <YOLOX_CHECKPOINT>

For example:

python tools/trt.py -n yolox-s -c your_ckpt.pth

<YOLOX_MODEL_NAME> can be: yolox-nano, yolox-tiny. yolox-s, yolox-m, yolox-l, yolox-x.

If you want to convert your customized model, use the flag -f to specify you exp file:

python tools/trt.py -f <YOLOX_EXP_FILE> -c <YOLOX_CHECKPOINT>

For example:

python tools/trt.py -f /path/to/your/yolox/exps/yolox_s.py -c your_ckpt.pth

yolox_s.py can be any exp file modified by you.

The converted model and the serialized engine file (for C++ demo) will be saved on your experiment output dir.

The TensorRT python demo is merged on our pytorch demo file, so you can run the pytorch demo command with --trt.

python tools/demo.py image -n yolox-s --trt --save_result

or

python tools/demo.py image -f exps/default/yolox_s.py --trt --save_result