在 C++ 中使用 Torch-TensorRT¶

如果您还没有，请按照安装中的说明获取库的压缩包。

在 C++ 中使用 Torch-TensorRT¶

Torch-TensorRT C++ API 接受 TorchScript 模块（从 torch.jit.script 或 torch.jit.trace 生成）作为输入，并返回一个 Torchscript 模块（使用 TensorRT 优化）。这要求用户事先使用 Pytorch（在 python 中）生成 torchscript 模块。请参考在 Python 中创建 TorchScript 模块部分以生成 torchscript 图表。

[Torch-TensorRT 快速入门] 使用 `torchtrtc` 编译 TorchScript 模块¶

使用 Torch-TensorRT 并检查您的模型是否可以在没有额外工作的情况下得到支持的简单方法是将其通过 torchtrtc 运行，它支持编译器从命令行提供的几乎所有功能，包括训练后量化（给定先前创建的校准缓存）。例如，我们可以通过设置我们首选的操作精度和输入大小来编译我们的 lenet 模型。这个新的 TorchScript 文件可以加载到 Python 中（注意：在加载这些编译后的模块之前，您需要 import torch_tensorrt，因为编译器扩展了 PyTorch 的反序列化器和运行时来执行编译后的模块）。

❯ torchtrtc -p f16 lenet_scripted.ts trt_lenet_scripted.ts "(1,1,32,32)"

❯ python3
Python 3.6.9 (default, Apr 18 2020, 01:56:04)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_tensorrt
>>> ts_model = torch.jit.load(“trt_lenet_scripted.ts”)
>>> ts_model(torch.randn((1,1,32,32)).to(“cuda”).half())

您可以在这里了解有关 torchtrtc 用法的更多信息：torchtrtc

在 C++ 中使用 TorchScript¶

如果我们正在开发一个用 C++ 部署的应用程序，我们可以使用 torch.jit.save 保存我们跟踪的或脚本化的模块，它将 TorchScript 代码、权重和其他信息序列化到一个包中。这也标志着我们对 Python 的依赖关系结束。

torch_script_module.save("lenet.jit.pt")

从这里我们可以现在在 C++ 中加载我们的 TorchScript 模块

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
    torch::jit::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load("<PATH TO SAVED TS MOD>");
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }

    std::cout << "ok\n";

如果您愿意，您可以使用 PyTorch / LibTorch 在 C++ 中进行完整的训练和推理，甚至可以定义 C++ 中的模块，并访问支持 PyTorch 的相同强大的张量库。（更多信息：https://pytorch.ac.cn/cppdocs/）。例如，我们可以使用我们的 LeNet 模块进行推理，如下所示

mod.eval();
torch::Tensor in = torch::randn({1, 1, 32, 32});
auto out = mod.forward(in);

并在 GPU 上运行

mod.eval();
mod.to(torch::kCUDA);
torch::Tensor in = torch::randn({1, 1, 32, 32}, torch::kCUDA);
auto out = mod.forward(in);

如您所见，它与 Python API 非常相似。当您调用 forward 方法时，您将调用 PyTorch JIT 编译器，它将优化并运行您的 TorchScript 代码。

使用 Torch-TensorRT 在 C++ 中编译¶

我们也处于可以使用 Torch-TensorRT 编译和优化我们的模块的阶段，但与 JIT 方式不同的是，我们必须在提前 (AOT) 进行，即在我们开始实际推理工作之前，因为它需要花费一些时间来优化模块，在每次运行模块时或甚至在第一次运行时这样做都没有意义。

在加载了我们的模块后，我们可以将其馈送到 Torch-TensorRT 编译器。这样做时，我们必须提供有关预期输入大小的信息，并配置任何其他设置。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();

    auto in = torch::randn({1, 1, 32, 32}, {torch::kCUDA});
    auto trt_mod = torch_tensorrt::CompileGraph(mod, std::vector<torch_tensorrt::CompileSpec::InputRange>{{in.sizes()}});
    auto out = trt_mod.forward({in});

就这样！现在图形主要不是使用 JIT 编译器，而是使用 TensorRT 运行（尽管我们使用 JIT 运行时执行图形）。

我们还可以设置诸如操作精度之类的设置以 FP16 运行。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();

    auto in = torch::randn({1, 1, 32, 32}, {torch::kCUDA}).to(torch::kHALF);
    auto input_sizes = std::vector<torch_tensorrt::CompileSpec::InputRange>({in.sizes()});
    torch_tensorrt::CompileSpec info(input_sizes);
    info.enable_precisions.insert(torch::kHALF);
    auto trt_mod = torch_tensorrt::CompileGraph(mod, info);
    auto out = trt_mod.forward({in});

现在我们正在 FP16 精度下运行模块。然后，您可以保存模块以供以后加载。

trt_mod.save("<PATH TO SAVED TRT/TS MOD>")

Torch-TensorRT 编译的 TorchScript 模块的加载方式与普通 TorchScript 模块相同。确保您的部署应用程序链接到 libtorchtrt.so

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"

int main(int argc, const char* argv[]) {
    torch::jit::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load("<PATH TO SAVED TRT/TS MOD>");
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }

    torch::Tensor in = torch::randn({1, 1, 32, 32}, torch::kCUDA);
    auto out = mod.forward(in);

    std::cout << "ok\n";
}

如果您想保存 Torch-TensorRT 生成的引擎以便在 TensorRT 应用程序中使用，可以使用 ConvertGraphToTRTEngine API。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();

    auto in = torch::randn({1, 1, 32, 32}, {torch::kCUDA}).to(torch::kHALF);
    auto input_sizes = std::vector<torch_tensorrt::CompileSpec::InputRange>({in.sizes()});
    torch_tensorrt::CompileSpec info(input_sizes);
    info.enabled_precisions.insert(torch::kHALF);
    auto trt_mod = torch_tensorrt::ConvertGraphToTRTEngine(mod, "forward", info);
    std::ofstream out("/tmp/engine_converted_from_jit.trt");
    out << engine;
    out.close();

幕后¶

当向 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)

现在图形已从一组模块（每个模块管理自己的参数）转换到单个图形，其中参数内联到图形中，并且所有操作都已布置好。Torch-TensorRT 还执行了许多优化和映射，以便更轻松地将图形转换为 TensorRT。从这里，编译器可以通过遵循图形中的数据流来组装 TensorRT 引擎。

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

以下是编译完成后返回的图形

graph(%self_1 : __torch__.lenet, %input_0 : Tensor):
    %1 : ...trt.Engine = prim::GetAttr[name="lenet"](%self_1)
    %3 : Tensor[] = prim::ListConstruct(%input_0)
    %4 : Tensor[] = trt::execute_engine(%3, %1)
    %5 : Tensor = prim::ListUnpack(%4)
    return (%5)

您可以看到执行引擎的调用，在提取包含引擎的属性并构造输入列表之后，然后将张量返回给用户。

处理不支持的操作符¶

Torch-TensorRT 是一个新的库，PyTorch 操作符库非常庞大，因此编译器可能不支持某些操作符。您可以使用上面显示的组合技术来使模块完全支持 Torch-TensorRT，而那些不支持的模块则在部署应用程序中进行缝合，或者您可以为缺少的操作符注册转换器。

您可以使用 torch_tensorrt::CheckMethodOperatorSupport(const torch::jit::Module& module, std::string method_name) API 检查支持情况，而无需经历完整的编译管道，以查看哪些操作符不受支持。 torchtrtc 在开始编译之前会自动使用此方法检查模块，并会打印出不受支持的操作符列表。

注册自定义转换器¶

操作通过使用模块化转换器（一个函数，它接受 JIT 图形中的一个节点并生成 TensorRT 中的等效层或子图）映射到 TensorRT。Torch-TensorRT 附带一个转换器库，存储在一个注册表中，该注册表将根据正在解析的节点执行。例如，aten::relu(%input0.4) 指令将触发对它的 relu 转换器运行，从而在 TensorRT 图形中生成一个激活层。但是，由于此库并不完整，您可能需要编写自己的库才能使 Torch-TensorRT 支持您的模块。

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

例如，如果我们尝试使用不支持扁平化操作（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 来导入。我们首先创建一个自注册类 torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns() 的实例，该实例将在全局转换器注册表中注册转换器，并将像 aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor) 这样的函数模式与一个 lambda 关联起来，该 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;
        }
    });

int main() {
    ...

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

您可以在贡献者文档（编写转换器）中找到有关编写转换器所有细节的更多信息。如果您发现自己拥有大量的转换器实现库，请考虑将它们上游，欢迎 PR，这对社区来说将是一件好事。

在 C++ 中使用 Torch-TensorRT¶