英伟达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
的文件中具有优先权
- 需要从NVIDIA网站下载TensorRT和cuDNN的压缩包。
- 将这些文件放在目录中(
third_party/dist_dir/[x86_64-linux-gnu | aarch64-linux-gnu]
为此目的而存在目录) - 编译使用:
bazel build // :: libtrtorch --compilation_mode opt --distdir third_party / dist_dir / [x86_64-linux-gnu | aarch64-linux-gnu]
2.使用本地安装的cuDNN和TensorRT进行构建
如果发现错误并使用此方法进行编译,请在问题中进行披露(ldd
转储也可以)
- 开始编译之前,请在系统上安装TensorRT,CUDA和cuDNN。
- 在
WORKSPACE
评论中
- # Downloaded distributions to use with --distdir
- http_archive(
- name = "cudnn",
- urls = ["<URL>",],
- build_file = "@//third_party/cudnn/archive:BUILD",
- sha256 = "<TAR SHA256>",
- 10. strip_prefix = "cuda"
11. )
- 12.
13. http_archive(
- 14. name = "tensorrt",
- 15. urls = ["<URL>",],
- 16.
- 17. build_file = "@//third_party/tensorrt/archive:BUILD",
- 18. sha256 = "<TAR SHA256>",
- 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”
)
- 编译使用:
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结构