torch_tensorrt

PyTorch JIT 的即时编译 (AOT)

Torch-TensorRT 是一个针对 PyTorch/TorchScript 的编译器，通过 NVIDIA 的 TensorRT 深度学习优化器和运行时，针对 NVIDIA GPU。与 PyTorch 的即时 (JIT) 编译器不同，Torch-TensorRT 是一个即时编译 (AOT) 编译器，这意味着在您部署 TorchScript 代码之前，您需要通过一个显式的编译步骤将标准的 TorchScript 程序转换为针对 TensorRT 引擎的模块。Torch-TensorRT 作为 PyTorch 扩展运行，并编译无缝集成到 JIT 运行时的模块。使用优化后的图编译后，与运行 TorchScript 模块的感觉应该没有区别。您还可以在编译时访问 TensorRT 的配置套件，因此您可以指定模块的操作精度（FP32/FP16/INT8）和其他设置。

示例用法

import torch_tensorrt

...

trt_ts_module = torch_tensorrt.compile(torch_script_module,
    inputs = [example_tensor, # Provide example tensor for input shape or...
        torch_tensorrt.Input( # Specify input object with shape and dtype
            min_shape=[1, 3, 224, 224],
            opt_shape=[1, 3, 512, 512],
            max_shape=[1, 3, 1024, 1024],
            # For static size shape=[1, 3, 224, 224]
            dtype=torch.half) # Datatype of input tensor. Allowed options torch.(float|half|int8|int32|bool)
    ],
    enabled_precisions = {torch.half}, # Run with FP16)

result = trt_ts_module(input_data) # run inference
torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedded Torchscript

安装

对于 Linux x86_64 平台，Pytorch 库默认使用预 CXX11 ABI。因此，请使用 python3 setup.py install。

在 Jetson 平台上，NVIDIA 提供 预构建的 Pytorch 轮文件。这些轮文件使用 CXX11 ABI 构建。因此，在 Jetson 平台上，请使用 python3 setup.py install --jetpack-version 4.6 --use-cxx11-abi

内部结构

当提供已追踪的模块给 Torch-TensorRT 时，编译器将内部表示转换为类似这样的形式

graph(%input.2 : Tensor):
    %2 : Float(84, 10) = prim::Constant[value=<Tensor>]()
    %3 : Float(120, 84) = prim::Constant[value=<Tensor>]()
    %4 : Float(576, 120) = prim::Constant[value=<Tensor>]()
    %5 : int = prim::Constant[value=-1]() # x.py:25:0
    %6 : int[] = prim::Constant[value=annotate(List[int], [])]()
    %7 : int[] = prim::Constant[value=[2, 2]]()
    %8 : int[] = prim::Constant[value=[0, 0]]()
    %9 : int[] = prim::Constant[value=[1, 1]]()
    %10 : bool = prim::Constant[value=1]() # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %11 : int = prim::Constant[value=1]() # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %12 : bool = prim::Constant[value=0]() # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %self.classifer.fc3.bias : Float(10) = prim::Constant[value= 0.0464  0.0383  0.0678  0.0932  0.1045 -0.0805 -0.0435 -0.0818  0.0208 -0.0358 [ CUDAFloatType{10} ]]()
    %self.classifer.fc2.bias : Float(84) = prim::Constant[value=<Tensor>]()
    %self.classifer.fc1.bias : Float(120) = prim::Constant[value=<Tensor>]()
    %self.feat.conv2.weight : Float(16, 6, 3, 3) = prim::Constant[value=<Tensor>]()
    %self.feat.conv2.bias : Float(16) = prim::Constant[value=<Tensor>]()
    %self.feat.conv1.weight : Float(6, 1, 3, 3) = prim::Constant[value=<Tensor>]()
    %self.feat.conv1.bias : Float(6) = prim::Constant[value= 0.0530 -0.1691  0.2802  0.1502  0.1056 -0.1549 [ CUDAFloatType{6} ]]()
    %input0.4 : Tensor = aten::_convolution(%input.2, %self.feat.conv1.weight, %self.feat.conv1.bias, %9, %8, %9, %12, %8, %11, %12, %12, %10) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %input0.5 : Tensor = aten::relu(%input0.4) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %input1.2 : Tensor = aten::max_pool2d(%input0.5, %7, %6, %8, %9, %12) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %input0.6 : Tensor = aten::_convolution(%input1.2, %self.feat.conv2.weight, %self.feat.conv2.bias, %9, %8, %9, %12, %8, %11, %12, %12, %10) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %input2.1 : Tensor = aten::relu(%input0.6) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %x.1 : Tensor = aten::max_pool2d(%input2.1, %7, %6, %8, %9, %12) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %input.1 : Tensor = aten::flatten(%x.1, %11, %5) # x.py:25:0
    %27 : Tensor = aten::matmul(%input.1, %4)
    %28 : Tensor = trt::const(%self.classifer.fc1.bias)
    %29 : Tensor = aten::add_(%28, %27, %11)
    %input0.2 : Tensor = aten::relu(%29) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %31 : Tensor = aten::matmul(%input0.2, %3)
    %32 : Tensor = trt::const(%self.classifer.fc2.bias)
    %33 : Tensor = aten::add_(%32, %31, %11)
    %input1.1 : Tensor = aten::relu(%33) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %35 : Tensor = aten::matmul(%input1.1, %2)
    %36 : Tensor = trt::const(%self.classifer.fc3.bias)
    %37 : Tensor = aten::add_(%36, %35, %11)
    return (%37)
(CompileGraph)

图已经从类似您的 PyTorch 模块是模块集合的集合，每个模块管理自己的参数，转换为具有参数内联到图中并且所有操作展开的单个图。Torch-TensorRT 还执行了一系列优化和映射，以便将图更容易地转换为 TensorRT。从这里，编译器可以通过跟踪图中的数据流来组装 TensorRT 引擎。

当图构建阶段完成后，Torch-TensorRT 生成一个序列化的 TensorRT 引擎。从这里，根据 API，此引擎返回给用户或移动到图构建阶段。在这里，Torch-TensorRT 创建一个 JIT 模块来执行 TensorRT 引擎，该引擎将由 Torch-TensorRT 运行时实例化和管理。

这是编译完成后返回的图

graph(%self.1 : __torch__.___torch_mangle_10.LeNet_trt,
    %2 : Tensor):
    %1 : int = prim::Constant[value=94106001690080]()
    %3 : Tensor = trt::execute_engine(%1, %2)
    return (%3)
(AddEngineToGraph)

您可以看到执行引擎的调用，基于一个常量，即引擎的 ID，告诉 JIT 如何找到引擎以及将被输入到 TensorRT 的输入张量。引擎表示与运行正常 PyTorch 模块所执行的精确相同的计算，但优化以在您的 GPU 上运行。

Torch-TensorRT 通过生成与图中看到的指令相对应的层或子图来从 TorchScript 转换。转换器是用于将特定操作映射到 TensorRT 中的层或子图的代码的小模块。并非所有操作都受支持，但如果您需要实现一个，您可以在 C++ 中实现。

注册自定义转换器

通过模块化转换器将操作映射到TensorRT，这是一种函数，它从一个JIT图中的节点生成在TensorRT中的等效层或子图。Torch-TensorRT附带了一个存储在注册表中的转换器库，这将根据正在解析的节点执行。例如，一个aten::relu(%input0.4)指令将触发relu转换器在其上运行，产生TensorRT图中的一个激活层。但由于这个库不是详尽的，您可能需要编写自己的转换器以使Torch-TensorRT支持您的模块。

Torch-TensorRT发行版附带内部核心API头文件。因此，您可以访问转换器注册表并为所需的操作添加转换器。

例如，如果我们尝试编译一个不支持flatten操作（aten::flatten）的Torch-TensorRT构建的图，您可能会看到以下错误

terminate called after throwing an instance of 'torch_tensorrt::Error'
what():  [enforce fail at core/conversion/conversion.cpp:109] Expected converter to be true but got false
Unable to convert node: %input.1 : Tensor = aten::flatten(%x.1, %11, %5) # x.py:25:0 (conversion.AddLayer)
Schema: aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)
Converter for aten::flatten requested, but no such converter was found.
If you need a converter for this operator, you can try implementing one yourself
or request a converter: https://www.github.com/NVIDIA/Torch-TensorRT/issues

我们可以在我们的应用程序中为这个操作注册一个转换器。所有构建转换器所需工具都可以通过包含Torch-TensorRT/core/conversion/converters/converters.h来导入。我们首先创建一个实例self-registering的class torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns()，它将在全局转换器注册表中注册转换器，将函数模式如aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)与一个lambda函数关联，该函数将接受转换状态、要转换的节点/操作以及节点的所有输入，并作为副作用在TensorRT网络中产生一个新层。参数作为按模式列表中列出的顺序的TensorRT ITensors和Torch IValues的可检查联合体向量传递。

以下是一个aten::flatten转换器的实现，我们可以在我们的应用程序中使用。在转换器实现中，您可以完全访问Torch和TensorRT库。因此，例如，我们可以通过在PyTorch中运行操作而不是像我们在以下flatten转换器中那样自己实现完整计算来快速获取输出大小。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
#include "torch_tensorrt/core/conversion/converters/converters.h"

static auto flatten_converter = torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns()
    .pattern({
        "aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)",
        [](torch_tensorrt::core::conversion::ConversionCtx* ctx,
           const torch::jit::Node* n,
           torch_tensorrt::core::conversion::converters::args& args) -> bool {
            auto in = args[0].ITensor();
            auto start_dim = args[1].unwrapToInt();
            auto end_dim = args[2].unwrapToInt();
            auto in_shape = torch_tensorrt::core::util::toVec(in->getDimensions());
            auto out_shape = torch::flatten(torch::rand(in_shape), start_dim, end_dim).sizes();

            auto shuffle = ctx->net->addShuffle(*in);
            shuffle->setReshapeDimensions(torch_tensorrt::core::util::toDims(out_shape));
            shuffle->setName(torch_tensorrt::core::util::node_info(n).c_str());

            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
            return true;
        }
    });

要在Python中使用此转换器，建议使用PyTorch的C++ / CUDA扩展模板来将您的转换器库包装到一个.so文件中，然后您可以使用Python应用程序中的ctypes.CDLL()加载它。

有关编写转换器的所有详细信息的更多信息，请参阅贡献者文档（编写转换器）。如果您发现自己有一个大型转换器实现库，请考虑将其上游化，PR是受欢迎的，并且让社区受益将是非常好的。