用於ONNX的TensorRT后端
解析ONNX模型以使用TensorRT執行。
另請參閱TensorRT文檔。
有關最近更改的列表,請參見changelog。
支持的TensorRT版本
Master分支上的開發適用於具有完整維度和動態架構shape支持的TensorRT 7.2.1的最新版本。
對於TensorRT的早期版本,請參考其各自的分支。
完整尺寸+動態架構
在具有動態架構支持的全維度模式下,構建INetwork對象,需要調用以下API:
C ++
const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
builder->createNetworkV2(explicitBatch)
Python
import tensorrt
explicit_batch = 1 << (int)(tensorrt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
builder.create_network(explicit_batch)
有關這些API用法的示例,請參見:
支持的算子
當前支持的ONNX算子,可在算子支持列表中找到。
安裝
依存關系
Dependencies
編譯
對於在docker中進行構建,建議按照main(TensorRT存儲庫)[https://github.com/NVIDIA/TensorRT#setting-up-the-build-environment]中的說明使用和設置Docker容器,來構建onnx -tensorrt庫。
clone存儲庫后,可以通過運行以下命令,來構建解析器庫和可執行文件:
cd onnx-tensorrt
mkdir build && cd build
cmake .. -DTENSORRT_ROOT=<path_to_trt> && make -j
// Ensure that you update your LD_LIBRARY_PATH to pick up the location of the newly built library:
export LD_LIBRARY_PATH=$PWD:$LD_LIBRARY_PATH
可執行用法
可以使用onnx2trt
可執行文件,將ONNX模型轉換為序列化的TensorRT引擎:
onnx2trt my_model.onnx -o my_engine.trt
ONNX模型也可以轉換為人們可讀的文本:
onnx2trt my_model.onnx -t my_model.onnx.txt
ONNX模型也可以通過ONNX的優化庫(由dsandler添加)進行優化。要優化ONNX模型並輸出新的模型,可使用它-m
來指定輸出模型名稱,用-O
指定以分號分隔的優化過程列表,進行應用:
onnx2trt my_model.onnx -O "pass_1;pass_2;pass_3" -m my_model_optimized.onnx
通過運行以下命令,查看更多所有可用的優化遍歷:
onnx2trt -p
通過運行以下命令,查看更多用法信息:
onnx2trt -h
Python模塊
ONNX-TensorRT解析器的Python,綁定打包在隨附的.whl
文件中。用來安裝
python3 -m pip install <tensorrt_install_dir>/python/tensorrt-7.x.x.x-cp<python_ver>-none-linux_x86_64.whl
TensorRT 7.2.1支持ONNX版本1.6.0。通過以下方式安裝:
python3 -m pip install onnx==1.6.0
可以通過運行以下命令,來安裝ONNX-TensorRT后端:
python3 setup.py install
ONNX-TensorRT Python后端使用
可以在Python中如下使用ONNX的TensorRT后端:
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)
C ++庫用法
模型解析器庫libnvonnxparser.so,在頭文件中聲明了其C ++ API:
NvOnnxParser.h
測試
安裝后(或在Docker容器內),可以按以下方式運行ONNX后端測試:
僅實模型測試:
python onnx_backend_test.py OnnxBackendRealModelTest
所有測試:
python onnx_backend_test.py
可以使用-v
標志,使輸出更加詳細。
預訓練模型
在ONNX格式預先訓練模式可以在找到ONNX Model Zoo
https://github.com/onnx/onnx-tensorrt