Torch-TensorRT 詳解¶
Torch-TensorRT 是一個針對 NVIDIA GPU 的 PyTorch 模型編譯器，透過 TensorRT 模型最佳化 SDK。其目標是為 PyTorch 模型提供更好的推論效能，同時仍保持 PyTorch 的出色人體工學設計。
Dynamo 前端¶
Dynamo 前端是 Torch-TensorRT 的預設前端。它利用來自 PyTorch 的 dynamo 編譯器堆疊。
torch.compile (即時)¶
torch.compile 是一個 JIT 編譯器堆疊，因此，編譯會延遲到首次使用時。這表示當圖形中的條件改變時，圖形將自動重新編譯。這為使用者提供了最大的執行時彈性，但限制了關於序列化的選項。
在底層，torch.compile 會將它認為可以降低到 Torch-TensorRT 的子圖委派出去。Torch-TensorRT 進一步將這些圖形降低為僅由 核心 ATen 運算子 或適用於 TensorRT 加速的精選「高階運算」組成的運算。子圖會進一步分割成將在 PyTorch 中執行的組件和將進一步編譯為 TensorRT 的組件，這取決於對運算子的支援。然後，TensorRT 引擎會取代支援的區塊，並將混合子圖返回到 torch.compile 以在呼叫時執行。
接受的格式¶
torch.fx GraphModule (
torch.fx.GraphModule)PyTorch 模組 (
torch.nn.Module)
傳回值¶
在首次呼叫時觸發編譯的盒裝函式
torch_tensorrt.dynamo.compile (預先)¶
torch_tensorrt.dynamo.compile 是一個 AOT 編譯器，模型在顯式編譯階段中編譯。然後，這些編譯產物可以被序列化並在稍後重新載入。圖形會通過 torch.export.trace 系統被降低為由 核心 ATen 運算子 或適用於 TensoRT 加速的精選「高階運算」組成的圖形。子圖會進一步分割成將在 PyTorch 中執行的組件和將進一步編譯為 TensorRT 的組件，這取決於對運算子的支援。然後，TensorRT 引擎會取代支援的區塊，並將混合子圖封裝到可以序列化和重新載入的 ExportedProgram 中。
接受的格式¶
torch.export.ExportedProgram (
torch.export.ExportedProgram)torch.fx GraphModule (
torch.fx.GraphModule) (透過torch.export.export)PyTorch 模組 (
torch.nn.Module) (透過torch.export.export)
傳回值¶
torch.fx.GraphModule (可使用
torch.export.ExportedProgram序列化)
舊版前端¶
由於多年來 PyTorch 生態系統中出現了許多編譯器技術，Torch-TensorRT 有一些針對它們的舊版功能。
TorchScript (torch_tensorrt.ts.compile)¶
TorchScript 前端是 Torch-TensorRT 的原始預設前端，目標是 TorchScript 格式的模型。提供的圖形將被分割成支援和不支援的區塊。支援的區塊將被降低到 TensorRT，而不支援的區塊將保持不變，以使用 LibTorch 執行。產生的圖形會以 ScriptModule 的形式返回給使用者，可以使用 Torch-TensorRT PyTorch 執行階段擴充功能載入和儲存。
接受的格式¶
TorchScript 模組 (
torch.jit.ScriptModule)PyTorch 模組 (
torch.nn.Module) (透過torch.jit.script或torch.jit.trace)
傳回值¶
TorchScript 模組 (
torch.jit.ScriptModule)
FX 圖模組 (torch_tensorrt.fx.compile)¶
此前端幾乎完全被 Dynamo 前端取代，Dynamo 前端是透過 FX 前端提供的功能的超集。原始 FX 前端仍保留在程式碼庫中，以實現向後相容性。
接受的格式¶
torch.fx GraphModule (
torch.fx.GraphModule)PyTorch 模組 (
torch.nn.Module) (透過torch.fx.trace)
傳回值¶
torch.fx GraphModule (
torch.fx.GraphModule)
torch_tensorrt.compile¶
由於有許多不同的前端和支援的格式，我們提供了一個稱為 torch_tensorrt.compile 的便利層，讓使用者可以存取所有不同的編譯器選項。您可以透過設定 ir 選項，告訴 Torch-TensorRT 嘗試透過特定的中間表示法來降低提供的模型，從而在 torch_tensorrt.compile 中指定要使用的編譯器路徑。
ir 選項¶
torch_compile：使用torch.compile系統。立即傳回一個盒裝函式，該函式將在首次呼叫時編譯dynamo：透過torch.export/ torchdynamo 堆疊執行圖形。如果輸入模組是torch.nn.Module，則它必須是「可匯出追蹤」的，因為該模組將使用torch.export.export進行追蹤。傳回一個torch.fx.GraphModule，可以立即執行或透過torch.export.export或torch_tensorrt.save儲存torchscript或ts：透過 TorchScript 堆疊執行圖形。如果輸入模組是torch.nn.Module，則它必須是「可腳本化」的，因為該模組將使用torch.jit.script進行編譯。傳回一個torch.jit.ScriptModule，可以立即執行或透過torch.save或torch_tensorrt.save儲存fx：透過torch.fx堆疊執行圖形。如果輸入模組是torch.nn.Module，它將使用torch.fx.trace進行追蹤，並受其限制。
