运行时阶段

运行时阶段负责构建包含嵌入式 TensorRT 引擎的独立 TorchScript 图，并在调用这些引擎时充当运行时环境。主要接口接受一个序列化的 TensorRT 引擎。执行阶段将反序列化此引擎，并将其包装在一个类中；该类为每个引擎维护一个执行上下文，并包含有关其输入和输出的一些元数据，并且与 TorchScript 解释器兼容，以便它可以像其他 TorchScript IValues 一样被移动和使用。通过将引擎及其输入提供给 tensorrt::execute_engine 算子来运行引擎，该算子将接收引擎及其输入，并返回引擎执行的结果。

背景

PyTorch JIT 的运行时基于栈式机器，所有算子都从栈中弹出参数，将它们传递给算子的某个实现，然后将结果压回栈中。栈的实际元素是 torch::jit::IValues，与我们在转换阶段评估的类型相同（即抽象类型 torch::jit::Value 的具体实现）。

TensorRT 引擎执行器算子

加载 Torch-TensorRT 时，它会在 PyTorch JIT 算子库中注册一个名为 trt::execute_engine(Tensor[] inputs, __torch__.torch.classes.tensorrt.Engine engine) -> Tensor[] 的算子，该算子接受一个实例化后的引擎和输入列表。编译后的图将此引擎存储在一个属性中，以便它具有可移植性和可序列化性。调用该算子时，实例化后的引擎和输入张量会从运行时栈中弹出。这些输入被传递到一个通用引擎执行函数中，该函数将张量通过 TensorRT 引擎运行，并返回新的张量作为结果。这些张量被压入栈中，以便下一个算子（无论是什么）可以使用它们。

构建结果图

引擎反序列化并实例化后，编译器将构建一个图，该图在模块被调用时执行引擎。以下是一个示例

graph(%self_1 : __torch__.torchvision.models.resnet.___torch_mangle_4847.ResNet_trt,
  %input_0 : Tensor):
    %1 : __torch__.torch.classes.tensorrt.Engine = prim::GetAttr[name="__torch___torchvision_models_resnet____torch_mangle_4847_ResNet_trt_engine"](%self_1)
    %3 : Tensor[] = prim::ListConstruct(%input_0)
    %4 : Tensor[] = trt::execute_engine(%3, %1)
    %5 : Tensor = prim::ListUnpack(%4)
return (%5)

您可以在图中看到引擎属性以及 trt::execute_engine 算子，它接受输入张量列表和一个引擎作为输入，并产生一个输出张量列表作为返回值。当在模块上调用 forward 时，此图将被执行，从而运行 TensorRT 引擎。

在存在多个输出的情况下，编译后的图可能会将输出张量重新打包到一个元组 (Tuple) 中返回给用户。

graph(%self_1 : __torch__.PyTorch.Detection.SSD.src.model.SSD300_trt,
  %input_0 : Tensor):
    %1 : __torch__.torch.classes.tensorrt.Engine = prim::GetAttr[name="__torch___PyTorch_Detection_SSD_src_model_SSD300_trt_engine"](%self_1)
    %3 : Tensor[] = prim::ListConstruct(%input_0)
    %4 : Tensor[] = trt::execute_engine(%3, %1)
    %5 : Tensor, %6 : Tensor = prim::ListUnpack(%4)
    %7 : (Tensor, Tensor) = prim::TupleConstruct(%5, %6)
return (%7)

序列化和反序列化

嵌入在 TorchScript 图中的 TensorRT 引擎的序列化和反序列化由引擎的持有者类 (holder class) 和 TorchBind 处理。保存 TorchScript 模块时，pickler 将对 cuda 引擎执行序列化，并将序列化的引擎存储在创建的 zip 文件中。反序列化时，depickler 将使用序列化的引擎调用引擎持有者类的构造函数，以便它可以再次设置为可执行状态。

ABI 版本控制和序列化格式

Torch-TensorRT 程序是标准的 TorchScript，其中包含作为对象嵌入在图中的 TensorRT 引擎。因此，存在 TensorRT 引擎的序列化格式。Torch-TensorRT 序列化程序的格式通过“ABI”版本进行版本控制，该版本告诉运行时有关运行时兼容性的信息。

> 当前 ABI 版本为 3

该格式是一个序列化字符串向量。它们编码了以下信息

• 程序的 ABI 版本

• TRT 引擎的名称

• 设备信息：包括构建引擎的目标设备、SM 能力及其他设备信息。此信息在反序列化时用于选择运行引擎的正确设备

• 序列化的 TensorRT 引擎