运行时阶段¶

运行时阶段负责构建包含嵌入式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引擎。

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

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引擎的序列化和反序列化由引擎的持有者类和TorchBind处理。当保存TorchScript模块时，pickler将对cuda引擎运行序列化，并将序列化的引擎存储在创建的zip文件中。在反序列化时，depickler将使用序列化的引擎调用引擎持有者类的构造函数，以便可以再次设置以执行。

ABI版本控制和序列化格式¶

Torch-TensorRT 程序是标准的 TorchScript，其中 TensorRT 引擎作为嵌入图中的对象。因此，TensorRT 引擎有一个序列化格式。Torch-TensorRT 序列化程序的格式使用“ABI”版本进行版本控制，该版本告诉运行时关于运行时的兼容性。

> 当前ABI版本为3

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

程序的ABI版本
TRT引擎的名称
设备信息：包括引擎构建的目标设备、SM能力和其他设备信息。此信息在反序列化时用于选择运行引擎的正确设备。
序列化的TensorRT引擎