在 C++ 中使用 Torch-TensorRT¶
如果您还没有这样做，请按照安装中的说明获取库的压缩包。
在C++中使用Torch-TensorRT¶
Torch-TensorRT C++ API 接受 TorchScript 模块（由 torch.jit.script 或 torch.jit.trace 生成）作为输入，并返回一个使用 TensorRT 优化的 TorchScript 模块。然而，C++ API 不支持 Dynamo 编译工作流，但对于 FX 和 Dynamo 工作流，支持执行 torch.jit.trace 编译的 FX GraphModules。
请参考在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";
如果你想的话，你可以在C++中使用PyTorch / LibTorch进行完整的训练和推理，你甚至可以在C++中定义你的模块，并访问支持PyTorch的强大张量库。（更多信息：https://pytorch.org/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头文件。因此，您可以访问转换器注册表并为您需要的操作添加转换器。
例如，如果我们尝试使用不支持展平操作（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 Extension模板将您的转换器库包装成一个.so文件，您可以在Python应用程序中使用ctypes.CDLL()加载它。
您可以在贡献者文档（writing_converters）中找到有关编写转换器的所有详细信息。如果您发现自己有大量的转换器实现库，请考虑将它们上游化，欢迎提交PR，这对社区也将大有裨益。
