用於 torch.compile 的 TensorRT 後端¶
本指南介紹 Torch-TensorRT torch.compile 後端：一個深度學習編譯器，使用 TensorRT 加速各種模型的 JIT 風格工作流程。
主要特性¶
Torch-TensorRT torch.compile 後端的主要目標是結合 torch.compile API 的簡單性和 TensorRT 的效能，從而實現即時編譯 (Just-In-Time compilation) 工作流程。呼叫 torch.compile 後端非常簡單，只需匯入 torch_tensorrt 包並指定後端即可
import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False)
注意
使用者可以使用許多額外的自定義選項。這些選項將在本指南中進一步深入討論。
該後端可以處理各種複雜的模型結構，並提供一個易於使用的介面，以有效加速模型。此外，它還提供許多自定義選項，以確保編譯過程適合特定的用例。
可自定義設定¶
- class torch_tensorrt.dynamo.CompilationSettings(enabled_precisions: ~typing.Set[~torch_tensorrt._enums.dtype] = <factory>, debug: bool = False, workspace_size: int = 0, min_block_size: int = 5, torch_executed_ops: ~typing.Collection[~typing.Union[~typing.Callable[[...], ~typing.Any], 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_double: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, device: ~torch_tensorrt._Device.Device = <factory>, require_full_compilation: bool = False, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, engine_capability: ~torch_tensorrt._enums.EngineCapability = <factory>, num_avg_timing_iters: int = 1, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, dryrun: ~typing.Union[bool, str] = False, hardware_compatible: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', lazy_engine_init: bool = False, cache_built_engines: bool = False, reuse_cached_engines: bool = False, use_explicit_typing: bool = False, use_fp32_acc: bool = False, refit_identical_engine_weights: bool = False, strip_engine_weights: bool = False, immutable_weights: bool = True, enable_weight_streaming: bool = False, enable_cross_compile_for_windows: bool = False, tiling_optimization_level: str = 'none', l2_limit_for_tiling: int = -1, use_distributed_mode_trace: bool = False)[source]¶
Torch-TensorRT Dynamo 路徑的編譯設定
- 引數
enabled_precisions (Set[dpython:type]) – 可用的核心 dtype 精度
debug (bool) – 是否列印詳細的除錯資訊
workspace_size (python:int) – TRT 允許用於模組的工作空間大小 (0 為預設值)
min_block_size (python:int) – 每個 TRT 引擎塊的最小運算子數量
torch_executed_ops (Collection[Target]) – 無論轉換器是否覆蓋，都在 Torch 中執行的操作集合
pass_through_build_failures (bool) – 是否在 TRT 引擎構建失敗時繼續 (True) 或終止 (False)
max_aux_streams (Optional[python:int]) – 每個引擎允許的最大輔助 TRT 流數量
version_compatible (bool) – 為引擎 plan 檔案提供版本向前相容性
optimization_level (Optional[python:int]) – 構建器最佳化級別 0-5，級別越高意味著構建時間越長，會搜尋更多最佳化選項。TRT 預設為 3
use_python_runtime (Optional[bool]) – 是否嚴格使用 Python 執行時或 C++ 執行時。要根據 C++ 依賴項的存在情況自動選擇執行時（優先選擇 C++ 執行時，如果可用），請將此引數保留為 None
truncate_double (bool) – 是否將 float64 TRT 引擎輸入或權重截斷為 float32
use_fast_partitioner (bool) – 是否使用快速或全域性圖劃分系統
enable_experimental_decompositions (bool) – 是否啟用所有核心 aten 分解，還是僅啟用選定子集
device (Device) – 用於編譯模型的 GPU
require_full_compilation (bool) – 是否要求圖在 TensorRT 中完全編譯。僅適用於 ir=”dynamo”；對 torch.compile 路徑無效
assume_dynamic_shape_support (bool) – 將此設定為 True 會使轉換器同時支援動態和靜態形狀。預設值：False
disable_tf32 (bool) – 是否停用 TRT 層的 TF32 計算
sparse_weights (bool) – 是否允許構建器使用稀疏權重
engine_capability (trt.EngineCapability) – 將核心選擇限制為安全的 GPU 核心或安全的 DLA 核心
num_avg_timing_iters (python:int) – 用於選擇核心的平均計時迭代次數
dla_sram_size (python:int) – DLA 用於層內通訊的快速軟體管理 RAM。
dla_local_dram_size (python:int) – DLA 用於在操作間共享中間張量資料的主機 DRAM
dla_global_dram_size (python:int) – DLA 用於儲存權重和元資料以供執行的主機 DRAM
dryrun (Union[bool, str]) – 切換“空執行”(Dryrun) 模式，該模式會執行除轉換為 TRT 引擎之外的所有分割槽過程。列印詳細的圖結構和分割槽性質日誌。如果指定字串路徑，則可以選擇將輸出儲存到檔案
hardware_compatible (bool) – 構建與構建引擎所用 GPU 架構之外的 GPU 架構相容的 TensorRT 引擎（目前適用於 NVIDIA Ampere 及更新的架構）
timing_cache_path (str) – 計時快取的路徑（如果存在）或編譯後將儲存計時快取的位置
cache_built_engines (bool) – 是否將編譯好的 TRT 引擎儲存到儲存
reuse_cached_engines (bool) – 是否從儲存載入快取的 TRT 引擎
use_strong_typing (bool) – 此標誌啟用 TensorRT 編譯中的強型別檢查，該檢查會遵循 PyTorch 模型中設定的精度。當用戶具有混合精度圖時，這非常有用。
use_fp32_acc (bool) – 此選項在 matmul 層周圍插入轉換為 FP32 的節點，TensorRT 確保 matmul 的累積在 FP32 中進行。僅當在 enabled_precisions 中配置了 FP16 精度時才使用此選項。
refit_identical_engine_weights (bool) – 是否使用相同權重重新適配引擎
strip_engine_weights (bool) – 是否剝離引擎權重
immutable_weights (bool) – 構建不可重新適配的引擎。這對於一些不可重新適配的層很有用。如果此引數設定為 true，則 strip_engine_weights 和 refit_identical_engine_weights 將被忽略
enable_weight_streaming (bool) – 啟用權重流式載入。
enable_cross_compile_for_windows (bool) – 預設情況下為 False，意味著 TensorRT 引擎只能在構建它的同一平臺上執行。設定為 True 將啟用跨平臺相容性，允許引擎在 Linux 上構建並在 Windows 上執行
tiling_optimization_level (str) – 切片 (Tiling) 策略的最佳化級別。級別越高，TensorRT 會花費更多時間搜尋更好的切片策略。我們目前支援 [“none”, “fast”, “moderate”, “full”]。
l2_limit_for_tiling (python:int) – 切片最佳化的目標 L2 快取使用限制（以位元組為單位）（預設值為 -1，表示無限制）。
use_distributed_mode_trace (bool) – 使用 aot_autograd 跟蹤圖。當分散式模型中存在 DTensors 或分散式張量時，此選項被啟用
自定義設定使用¶
import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
options={"truncate_long_and_double": True,
"enabled_precisions": {torch.float, 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 返回的最佳化模型對於模型基準測試很有用，因為它可以自動處理編譯上下文的變化或可能需要重新編譯的不同輸入。當對不同分佈、批處理大小或其他標準的輸入進行基準測試時，這可以節省時間。
運算子覆蓋範圍¶
編譯也是確定特定模型運算子覆蓋範圍的有用工具。例如，以下編譯命令將顯示每個圖的運算子覆蓋範圍，但不會編譯模型——有效地提供了一種“空執行”(dryrun) 機制
import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
options={"debug": True,
"min_block_size": float("inf"),})
如果模型的關鍵運算子不受支援，請參閱 dynamo_conversion 貢獻您自己的轉換器，或在此處提交問題：https://github.com/pytorch/TensorRT/issues。
序列化的可行性¶
編譯還可以幫助展示圖中斷 (graph breaks) 以及特定模型的序列化可行性。例如，如果一個模型沒有圖中斷，並且使用 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 處理這兩種情況下的保護措施 (guarding)，並確定何時需要重新編譯。
