在 C++ 中使用 Torch-TensorRT¶
如果您还没有，请按照安装中的说明获取库的 tarball 文件。
在 C++ 中使用 Torch-TensorRT¶
Torch-TensorRT C++ API 接受 TorchScript 模块（由 torch.jit.script 或 torch.jit.trace 生成）作为输入，并返回一个 Torchscript 模块（使用 TensorRT 优化）。C++ API 不支持 Dynamo 编译工作流，但支持执行通过 torch.jit.trace 编译的 FX GraphModule，以用于 FX 和 Dynamo 工作流。
请参考在 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 代码。
在 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 头文件。因此，您可以访问转换器注册表并为您需要的操作添加一个转换器。
例如，如果我们尝试用一个不支持 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) 与一个 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() 加载它。
您可以在贡献者文档（writing_converters）中找到更多关于编写转换器所有细节的信息。如果您发现自己拥有一个庞大的转换器实现库，请考虑将它们贡献给上游，欢迎提交 PR，这也会让社区受益。
