在 C++ 中使用 Torch-TensorRT¶
如果您尚未获取库的 tarball，请按照安装中的说明进行操作
在 C++ 中使用 Torch-TensorRT¶
Torch-TensorRT C++ API 接受 TorchScript 模块（从 torch.jit.script 或 torch.jit.trace 生成）作为输入，并返回 Torchscript 模块（使用 TensorRT 优化）。C++ API 中将不支持 Dynamo 编译工作流程，但是，对于 FX 和 Dynamo 工作流程，支持执行 torch.jit.trace'd 编译的 FX GraphModules。
请参阅在 Python 中创建 TorchScript 模块部分，以生成 torchscript 图。
[Torch-TensorRT 快速入门] 使用 torchtrtc 编译 TorchScript 模块¶
开始使用 Torch-TensorRT 并检查您的模型是否可以在无需额外工作的情况下获得支持的简单方法是通过 torchtrtc 运行它，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 编译和优化我们的模块的地步，但是我们必须提前 (AOT) 进行，而不是以 JIT 方式进行，即在我们开始进行实际的推理工作之前，因为它需要一些时间来优化模块，所以每次运行模块甚至第一次运行模块时都这样做是没有意义的。
加载我们的模块后，我们可以将其馈送到 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，这将对社区有很大的好处。
