英伟达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. build_file = "@//third_party/cudnn/archive:BUILD",
6. sha256 = "<TAR SHA256>",
7. 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” 

）

1. 编译使用：

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结构