英偉達TRTTorch

英偉達TRTTorch

PyTorch JIT的提前（AOT）編譯Ahead of Time (AOT) compiling for PyTorch JIT

TRTorch是PyTorch / TorchScript的編譯器，通過NVIDIA針對NVIDIA GPU的TensorRT深度學習優化器和運行時runtime。與PyTorch的即時（JIT）編譯器不同，TRTorch是一種提前（AOT）編譯器，這意味着在部署TorchScript代碼之前，需要執行顯式的編譯步驟，以TensorRT引擎為目標，將標准的TorchScript程序轉換為模塊。TRTorch充當PyTorch擴展，編譯與JIT運行時runtime無縫集成的模塊。使用優化圖進行編譯后，應該感覺與運行TorchScript模塊沒有什么不同。還可以在編譯時訪問TensorRT的配置套件，因此可以為模塊指定算子精度（FP32 / FP16 / INT8）和其他設置。

用法示例

C ++

#include "torch/script.h"

#include "trtorch/trtorch.h"

 

...

auto compile_settings = trtorch::CompileSpec(dims);

// FP16 execution

compile_settings.op_precision = torch::kFloat;

// Compile module

auto trt_mod = trtorch::CompileGraph(ts_mod, compile_settings);

// Run like normal

auto results = trt_mod.forward({in_tensor});

// Save module for later

trt_mod.save("trt_torchscript_module.ts");

...

python

import trtorch

 

...

compile_settings = {

    "input_shapes": [

        {

            "min": [1, 3, 224, 224],

            "opt": [1, 3, 512, 512],

            "max": [1, 3, 1024, 1024]

        }, # For static size [1, 3, 224, 224]

    ],

    "op_precision": torch.half # Run with FP16

}

 

trt_ts_module = trtorch.compile(torch_script_module, compile_settings)

 

input_data = input_data.half()

result = trt_ts_module(input_data)

torch.jit.save(trt_ts_module, "trt_torchscript_module.ts")

以較低的精度運行時runtime的注意事項：

• 使用compile_spec.op_precision設置精度
• 編譯之前，模塊應使用FP32（FP16可支持半張量模型）
• 在FP16中，僅應將輸入張量轉換為FP16，其他精度使用FP32

平台支援

 

 

 依存關系

這些是以下用於驗證測試用例的依賴項。TRTorch可以與其他版本一起使用，但不能保證測試能夠通過。

• Bazel 3.7.0
• Libtorch 1.7.1（使用CUDA 11.0構建）
• CUDA 11.0
• cuDNN 8
• TensorRT 7.2.1.6

預構建的二進制文件和Wheel文件

發布：https : //github.com/NVIDIA/TRTorch/releases

編譯TRTorch

安裝依賴項

0.安裝Install Bazel

如果沒有安裝bazel，最簡單的方法是使用選擇https://github.com/bazelbuild/bazelisk的方法來安裝bazelisk

否則，可以按照以下說明安裝二進制文件https://docs.bazel.build/versions/master/install.html

最后，如果需要從源代碼進行編譯（例如aarch64，直到bazel為該體系結構分發二進制文件），則可以使用以下說明

export BAZEL_VERSION=<VERSION>

mkdir bazel

cd bazel

curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-dist.zip

unzip bazel-$BAZEL_VERSION-dist.zip

bash ./compile.sh

需要先在系統上安裝CUDA，LibTorch會自動被bazel提起，然后有兩個選擇。

1.使用cuDNN和TensorRT tarball發行版進行構建

推薦這樣做是為了構建TRTorch，並確保任何錯誤不是由版本問題引起的

確保在運行TRTorch時，這些版本的庫在$LD_LIBRARY_PATH的文件中具有優先權 

1. 需要從NVIDIA網站下載TensorRT和cuDNN的壓縮包。
• https://developer.nvidia.com/cudnn
• https://developer.nvidia.com/tensorrt
2. 將這些文件放在目錄中（third_party/dist_dir/[x86_64-linux-gnu | aarch64-linux-gnu]為此目的而存在目錄）
3. 編譯使用：

bazel build // :: libtrtorch --compilation_mode opt --distdir third_party / dist_dir / [x86_64-linux-gnu | aarch64-linux-gnu]

2.使用本地安裝的cuDNN和TensorRT進行構建

如果發現錯誤並使用此方法進行編譯，請在問題中進行披露（ldd轉儲也可以）

1. 開始編譯之前，請在系統上安裝TensorRT，CUDA和cuDNN。
2. 在WORKSPACE評論中

1. # Downloaded distributions to use with --distdir
2. http_archive(
3.     name = "cudnn",
4.     urls = ["<URL>",],
5.  
6.     build_file = "@//third_party/cudnn/archive:BUILD",
7.     sha256 = "<TAR SHA256>",
8. 10.     strip_prefix = "cuda"

11. )

1. 12.  

13. http_archive(

1. 14.     name = "tensorrt",
2. 15.     urls = ["<URL>",],
3. 16.  
4. 17.     build_file = "@//third_party/tensorrt/archive:BUILD",
5. 18.     sha256 = "<TAR SHA256>",
6. 19.     strip_prefix = "TensorRT-<VERSION>"

20. )

和不加評論

＃本地安裝的依賴項

new_local_repository（

     name  =  “ cudnn”，

     path  =  “ / usr /”，

     build_file  =  “ @ // third_party / cudnn / local：BUILD”

）

new_local_repository（

    name  =  “ tensorrt”，

    path  =  “ / usr /”，

    build_file  =  “ @ // third_party / tensorrt / local：BUILD” 

）

3. 編譯使用：

bazel build // :: libtrtorch --compilation_mode選擇

調試版本

bazel build // :: libtrtorch --compilation_mode = dbg

NVIDIA Jetson AGX上的本機編譯

bazel build // :: libtrtorch --distdir third_party / dist_dir / aarch64-linux-gnu

注意：有關先決條件，請參閱安裝說明

然后可以在bazel-bin中找到包含包含文件和庫的tarball

在JIT圖上運行TRTorch

確保將LibTorch添加到的LD_LIBRARY_PATH
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/bazel-TRTorch/external/libtorch/lib

bazel run //cpp/trtorchexec -- $(realpath <PATH TO GRAPH>) <input-size>

編譯Python包

要為本地計算機編譯python軟件包，只需python3 setup.py install在//py目錄中運行即可。要為不同的python版本構建wheel文件，請先構建Dockerfile，//py然后運行以下命令

docker run -it -v$(pwd)/..:/workspace/TRTorch build_trtorch_wheel /bin/bash /workspace/TRTorch/py/build_whl.sh

Python編譯期望從上面使用基於tarball的編譯策略。

如何添加對新算子的支持...

在TRTorch？

支持新算子有兩種主要方法。可以從頭開始為op編寫一個轉換器並將其注冊到NodeConverterRegistry中，或者如果可以將op映射到一組已經具有轉換器的op上，則可以編寫圖形重寫過程，將新的op替換為等效的子圖支持的算子。首選使用圖形重寫，因為這樣就不需要維護大型的op轉換器庫。還要查看問題中的各種算子支持跟蹤器，以獲取有關各種算子的支持狀態的信息。

如何申請

Node Converter Registry不在頂級API中公開，而是在tarball附帶的內部header中公開。

可以使用NodeConverterRegistry應用程序內部為算子注冊一個轉換器。

Structure of the repo結構