Torch-TensorRT torch.compile 後端¶

本指南介紹 Torch-TensorRT torch.compile 後端：一種深度學習編譯器，它使用 TensorRT 來加速各種模型的 JIT 風格工作流程。

主要功能¶

Torch-TensorRT torch.compile 後端的主要目標是透過結合 torch.compile API 的簡單性和 TensorRT 的效能來實現即時編譯工作流程。呼叫 torch.compile 後端就像匯入 torch_tensorrt 套件並指定後端一樣簡單

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False)

注意

使用者可以使用許多其他自訂選項。本指南將更深入地討論這些選項。

後端可以處理各種具有挑戰性的模型結構，並提供易於使用的介面來有效加速模型。此外，它還具有許多自訂選項，以確保編譯過程適合特定的用例。

可自訂設定¶

class torch_tensorrt.dynamo.CompilationSettings(precision: torch.dtype = torch.float32, debug: bool = False, workspace_size: int = 0, min_block_size: int = 5, torch_executed_ops: typing.Set[str] = <factory>, pass_through_build_failures: bool = False, max_aux_streams: typing.Optional[int] = None, version_compatible: bool = False, optimization_level: typing.Optional[int] = None, use_python_runtime: typing.Optional[bool] = False, truncate_long_and_double: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, device: torch_tensorrt._Device.Device = <factory>, require_full_compilation: bool = False)[來源]¶

Torch-TensorRT Dynamo 路徑的編譯設定

參數

precision (torch.dpython:type) – 模型層精度
debug (bool) – 是否列印出詳細的偵錯資訊
workspace_size (python:int) – 允許 TRT 用於模組的工作區 (0 為預設值)
min_block_size (python:int) – 每個 TRT 引擎區塊的最小運算子數
torch_executed_ops (Sequence[str]) – 在 Torch 中執行的操作序列，不論轉換器覆蓋範圍
pass_through_build_failures (bool) – 是否在 TRT 引擎建置錯誤時失敗 (True) 或不失敗 (False)
max_aux_streams (Optional[python:int]) – 每個引擎允許的最大輔助 TRT 串流數
version_compatible (bool) – 為引擎計畫檔案提供版本向前相容性
optimization_level (Optional[python:int]) – 建置器最佳化 0-5，更高的級別意味著更長的建置時間，尋找更多最佳化選項。TRT 預設為 3
use_python_runtime (Optional[bool]) – 是否嚴格使用 Python 執行時間或 C++ 執行時間。若要根據 C++ 依賴項的存在自動選擇執行時間 (優先選擇 C++ 執行時間，如果有的話)，請將參數保留為 None
truncate_long_and_double (bool) – 是否將 int64/float64 TRT 引擎輸入或權重截斷為 int32/float32
use_fast_partitioner (bool) – 是否使用快速或全域圖形分割系統
enable_experimental_decompositions (bool) – 是否啟用所有核心 aten 分解，還是僅啟用其中選定的子集
device (Device) – 用於編譯模型的 GPU
require_full_compilation (bool) – 是否要求圖形在 TensorRT 中完全編譯。僅適用於 ir=”dynamo”；對 torch.compile 路徑沒有影響

自訂設定用法¶

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
                                options={"truncate_long_and_double": True,
                                         "precision": torch.half,
                                         "debug": True,
                                         "min_block_size": 2,
                                         "torch_executed_ops": {"torch.ops.aten.sub.Tensor"},
                                         "optimization_level": 4,
                                         "use_python_runtime": False,})

注意

量化/INT8 支援預計在未來版本中提供；目前，我們支援 FP16 和 FP32 精度層。

編譯¶

編譯是透過將輸入傳遞給模型來觸發的，因此

import torch_tensorrt
...
# Causes model compilation to occur
first_outputs = optimized_model(*inputs)

# Subsequent inference runs with the same, or similar inputs will not cause recompilation
# For a full discussion of this, see "Recompilation Conditions" below
second_outputs = optimized_model(*inputs)

編譯後¶

編譯後的物件可用於 Python 會話中的推論，並將根據以下詳細說明的重新編譯條件進行重新編譯。除了通用的推論之外，編譯過程還是一個有用的工具，可用於確定模型效能、當前的運算子覆蓋範圍和序列化可行性。以下將詳細介紹這些要點。

模型效能¶

torch.compile 返回的最佳化模型可用於模型基準測試，因為它可以自動處理編譯環境中的更改，或可能需要重新編譯的不同輸入。當對不同分佈、批次大小或其他標準的輸入進行基準測試時，這可以節省時間。

運算子覆蓋範圍¶

編譯也是一個有用的工具，可用於確定特定模型的運算子覆蓋範圍。例如，以下編譯命令將顯示每個圖形的運算子覆蓋範圍，但不會編譯模型 - 有效地提供了一種「空執行」機制

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
                                options={"debug": True,
                                         "min_block_size": float("inf"),})

如果您的模型的關鍵運算子不受支援，請參閱 Dynamo 轉換器以貢獻您自己的轉換器，或在此處提交問題：https://github.com/pytorch/TensorRT/issues。

序列化可行性¶

編譯也有助於演示圖形斷點和特定模型序列化之可行性。例如，如果模型沒有圖形斷點，並且使用 Torch-TensorRT 後端成功編譯，則該模型應該可以透過 torch_tensorrt Dynamo IR 進行編譯和序列化，如使用 Torch-TensorRT 的動態形狀中所述。若要判斷模型中圖形斷點的數量，torch._dynamo.explain 函數非常有用。

import torch
import torch_tensorrt
...
explanation = torch._dynamo.explain(model)(*inputs)
print(f"Graph breaks: {explanation.graph_break_count}")
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False, options={"truncate_long_and_double": True})

動態形狀支援¶

Torch-TensorRT torch.compile 後端目前需要針對遇到的每個新批次大小重新編譯，建議在使用此後端進行編譯時使用 dynamic=False 參數。計劃在未來版本中提供完整的動態形狀支援。

重新編譯條件¶

編譯模型後，後續具有相同形狀和資料類型的推論輸入，以相同方式遍歷圖形，將不需要重新編譯。此外，每個新的重新編譯都將在 Python 工作階段期間快取。例如，如果將批次大小為 4 和 8 的輸入提供給模型，導致兩次重新編譯，則在同一個工作階段內的推論過程中，未來使用這些批次大小的輸入將不需要進一步重新編譯。計劃在未來版本中支援引擎快取序列化。

重新編譯通常由以下兩個事件之一觸發：遇到不同大小的輸入或以不同方式遍歷模型程式碼的輸入。當模型程式碼包含條件邏輯、複雜迴圈或資料相依形狀時，可能會發生後一種情況。torch.compile 會處理這兩種情況下的防護，並確定何時需要重新編譯。