用於ONNX的TensorRT後端

wujianming_110117發表於2020-12-21

用於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用法的示例,請參見:

sampleONNXMNIST
sampleDynamicReshape

支援的運算元

當前支援的ONNX運算元,可在運算元支援列表中找到。

安裝

依存關係

Dependencies

Protobuf>= 3.0.x
TensorRT 7.2.1
TensorRT 7.2.1 open source libaries (master branch)

編譯

對於在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= P W D : PWD: 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

相關文章