在 C++ 中使用 Torch-TensorRT¶
如果尚未完成，请按照 安装 中的说明获取库的 tarball
在 C++ 中使用 Torch-TensorRT¶
Torch-TensorRT C++ API 接受 TorchScript 模块（由 torch.jit.script 或 torch.jit.trace 生成）作为输入，并返回一个 Torchscript 模块（使用 TensorRT 优化）。Dynamo 编译工作流将不受 C++ API 支持。但是，FX 和 Dynamo 工作流支持执行 torch.jit.trace 编译的 FX GraphModules。
要生成 torchscript 图，请参阅 在 Python 中创建 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 代码。
在 C++ 中使用 Torch-TensorRT 进行编译¶
我们也可以在这一点上使用 Torch-TensorRT 编译和优化我们的模块，但与 JIT 方式不同，我们必须提前（AOT）完成，即在开始实际推理工作之前，因为它需要一些时间来优化模块，所以每次运行模块或第一次运行时这样做都没有意义。
加载模块后，我们可以将其输入 Torch-TensorRT 编译器。在这样做时，我们必须提供有关预期输入大小的一些信息，并配置任何其他设置。
#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...
mod.to(at::kCUDA);
mod.eval();
std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
torch_tensorrt::ts::CompileSpec cfg(inputs);
auto trt_mod = torch_tensorrt::ts::compile(mod, cfg);
auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA});
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, 3, 224, 224}, {torch::kCUDA}).to(torch::kHALF);
std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
torch_tensorrt::ts::CompileSpec cfg(inputs);
cfg.enable_precisions.insert(torch::kHALF);
auto trt_mod = torch_tensorrt::ts::compile(mod, cfg);
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, 3, 224, 224}, {torch::kCUDA}).to(torch::kHALF);
std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
torch_tensorrt::ts::CompileSpec cfg(inputs);
cfg.enabled_precisions.insert(torch::kHALF);
auto trt_mod = torch_tensorrt::ts::convert_method_to_trt_engine(mod, "forward", cfg);
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.classifier.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.classifier.fc2.bias : Float(84) = prim::Constant[value=<Tensor>]()
%self.classifier.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.classifier.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.classifier.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.classifier.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在开始编译之前会自动检查此方法的模块，并打印出不支持的运算符列表。
注册自定义转换器¶
运算符通过使用模块化转换器映射到 TensorRT，转换器是一个函数，它接收 JIT 图中的一个节点并生成 TensorRT 中等效的层或子图。Torch-TensorRT 附带一个存储在注册表中的转换器库，该库将根据正在解析的节点执行。例如，aten::relu(%input0.4) 指令将触发 relu 转换器在其上运行，从而在 TensorRT 图中生成一个激活层。但由于此库并非详尽无遗，您可能需要编写自己的转换器才能让 Torch-TensorRT 支持您的模块。
Torch-TensorRT 发行版附带了内部核心 API 头文件。因此，您可以访问转换器注册表并为您需要的 op 添加转换器。
例如，如果我们尝试使用不支持 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 来导入。我们首先创建一个自注册类 torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns() 的实例，该类会将转换器注册到全局转换器注册表中，将类似于 aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor) 的函数 schema 与一个 lambda 关联起来，该 lambda 将接收转换状态、需要转换的节点/操作以及节点的所有输入，并作为副作用生成 TensorRT 网络中的一个新层。参数按 schema 中列出的顺序作为 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() 加载该文件。
您可以在贡献者文档（writing_converters）中找到有关编写转换器所有细节的更多信息。如果您发现自己拥有大量转换器实现，请考虑将它们合并到主线中，PR 是受欢迎的，这将使社区受益。
