The ONNX Runtime inference for yolort
, both CPU and GPU are supported.
- ONNX Runtime 1.7+
- OpenCV
- CUDA [Optional]
We didn't impose too strong restrictions on the versions of dependencies.
The ONNX model exported by yolort differs from other pipeline in the following three ways.
- We embed the pre-processing into the graph (mainly composed of
letterbox
). and the exported model expects aTensor[C, H, W]
, which is inRGB
channel and is rescaled to rangefloat32 [0-1]
. - We embed the post-processing into the model graph with
torchvision.ops.batched_nms
. So the outputs of the exported model are straightforwardboxes
,labels
andscores
fields of this image. - We adopt the dynamic shape mechanism to export the ONNX models.
-
Export your custom model to ONNX.
python tools/export_model.py --checkpoint_path {path/to/your/best.pt} --size_divisible 32/64
And then, you can find that a ONNX model ("best.onnx") have been generated in the directory of "best.pt". Set the
size_divisible
here according to your model, 32 for P5 ("yolov5s.pt" for instance) and 64 for P6 ("yolov5s6.pt" for instance). -
[Optional] Quick test with the ONNX Runtime Python interface.
from yolort.runtime import PredictorORT # Load the serialized ONNX model engine_path = "yolov5n6.onnx" device = "cpu" y_runtime = PredictorORT(engine_path, device=device) # Perform inference on an image file predictions = y_runtime.predict("bus.jpg")
-
Compile the source code.
cd deployment/onnxruntime mkdir build && cd build cmake .. -DONNXRUNTIME_DIR={path/to/your/ONNXRUNTIME/install/director} cmake --build .
-
Now, you can infer your own images.
./yolort_onnx --image ../../../test/assets/zidane.jpg --model_path ../../../notebooks/best.onnx --class_names ../../../notebooks/assets/coco.names [--gpu] # GPU switch, which is optional, and set False as default