執行階段

執行階段負責建構具有嵌入式 TensorRT 引擎的獨立 TorchScript 圖形，並在呼叫這些引擎時作為執行階段。主要介面接受序列化的 TensorRT 引擎。執行階段將會還原序列化此引擎，並將其包裝在一個類別中，該類別為每個引擎維護執行環境，以及關於其輸入和輸出的一些中繼資料，並且與 TorchScript 解譯器相容，以便可以像其他 TorchScript IValue 一樣移動和使用。引擎的運行方式是將引擎和輸入提供給 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 引擎的序列化和還原序列化由引擎和 TorchBind 的持有者類別處理。當儲存 TorchScript 模組時，pickler 將在 CUDA 引擎上運行序列化，並將序列化的引擎儲存在建立的 zip 檔案中。在還原序列化時，depickler 將使用序列化的引擎呼叫引擎持有者類別的建構函式，以便可以再次設定以進行執行。

ABI 版本控制和序列化格式

Torch-TensorRT 程式是標準 TorchScript，其中 TensorRT 引擎作為物件嵌入在圖形中。因此，TensorRT 引擎有一個序列化格式。Torch-TensorRT 序列化程式的格式使用「ABI」版本進行版本控制，該版本告知執行階段關於執行階段相容性的資訊。

> 目前的 ABI 版本為 3

該格式是一個序列化字串向量。它們編碼以下資訊

• 程式的 ABI 版本

• TRT 引擎的名稱

• 裝置資訊：包括引擎建構所在的目標裝置、SM 功能和其他裝置資訊。此資訊在還原序列化時用於選擇正確的裝置來運行引擎

• 序列化的 TensorRT 引擎