Parses ONNX models for execution with TensorRT.
See also the TensorRT documentation.
The TensorRT backend for ONNX can be used in Python as follows:
import onnx
import onnx_tensorrt.backend as backend
import numpy as np
model = onnx.load("/path/to/model.onnx")
engine = backend.prepare(model, device='CUDA:1')
input_data = np.random.random(size=(32, 3, 224, 224)).astype(np.float32)
output_data = engine.run(input_data)[0]
print(output_data)
print(output_data.shape)ONNX models can be converted to serialized TensorRT engines using the onnx2trt executable:
onnx2trt my_model.onnx -o my_engine.trt
ONNX models can also be converted to human-readable text:
onnx2trt my_model.onnx -t my_model.onnx.txt
See more usage information by running:
onnx2trt -h
The model parser library, libnvonnxparser.so, has a C++ API declared in this header:
NvOnnxParser.h
TensorRT engines built using this parser must use the plugin factory provided in libnvonnxparser_runtime.so, which has a C++ API declared in this header:
NvOnnxParserRuntime.h
Clone the code from GitHub.
git clone --recursive https://github.com/onnx/onnx-tensorrt.git
Suppose your TensorRT library is located at /opt/tensorrt. Build the onnx2trt executable and the libnvonnxparser* libraries using CMake:
mkdir build
cd build
cmake .. -DTENSORRT_ROOT=/opt/tensorrt
make -j8
sudo make install
TensorRT version 5.0 and up create Python wrappers using pybind11. No build is necessary. However, since we use already existing TensorRT binding, we need to ensure that PYTHONPATH is set appropriately. If you just install TensorRT, the weels for the pybind will be located at the same location other TensorRT libs.
For earlier versions of TensorRT, the Python wrappers are built using SWIG. Build the Python wrappers and modules by running:
python setup.py build
sudo python setup.py install
Build the onnx_tensorrt Docker image by running:
cp /path/to/TensorRT-3.0.*.tar.gz .
docker build -t onnx_tensorrt .
After installation (or inside the Docker container), ONNX backend tests can be run as follows:
Real model tests only:
python onnx_backend_test.py OnnxBackendRealModelTest
All tests:
python onnx_backend_test.py
You can use -v flag to make output more verbose.
Pre-trained Caffe2 models in ONNX format can be found at https://github.com/onnx/models