執行時階段

執行時階段負責構建包含嵌入式 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 引擎