torch_tensorrt.dynamo¶

函數¶

torch_tensorrt.dynamo.compile(exported_program: ExportedProgram, inputs: Optional[Sequence[Sequence[Any]]] = None, *, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, device: Optional[Union[Device, device, str]] = None, disable_tf32: bool = False, assume_dynamic_shape_support: bool = False, sparse_weights: bool = False, enabled_precisions: Union[Set[Union[dtype, dtype]], Tuple[Union[dtype, dtype]]] = {dtype.f32}, engine_capability: EngineCapability = EngineCapability.STANDARD, debug: bool = False, num_avg_timing_iters: int = 1, workspace_size: int = 0, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, truncate_double: bool = False, require_full_compilation: bool = False, min_block_size: int = 5, torch_executed_ops: Optional[Collection[Union[Callable[[...], Any], str]]] = None, torch_executed_modules: Optional[List[str]] = None, pass_through_build_failures: bool = False, max_aux_streams: Optional[int] = None, version_compatible: bool = False, optimization_level: Optional[int] = None, use_python_runtime: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, dryrun: bool = 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, engine_cache_dir: str = '/tmp/torch_tensorrt_engine_cache', engine_cache_size: int = 5368709120, custom_engine_cache: Optional[BaseEngineCache] = None, 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, **kwargs: Any) → GraphModule[source]¶

使用 TensorRT 為 NVIDIA GPU 編譯 ExportedProgram 模組

採用現有的 TorchScript 模組和一組設定來配置編譯器，並將方法轉換為呼叫同等 TensorRT 引擎的 JIT 圖形

專門轉換 TorchScript 模組的 forward 方法

參數

exported_program (torch.export.ExportedProgram) – 來源模組，在 torch.nn.Module 上執行 torch.export

inputs (Tuple[Any, ...]) –

模組輸入的輸入形狀、dtype 和記憶體配置規格清單。此引數為必要項。輸入大小可以指定為 torch 大小、元組或清單。dtype 可以使用 torch 資料類型或 torch_tensorrt 資料類型來指定，您可以使用 torch 裝置或 torch_tensorrt 裝置類型列舉來選擇裝置類型。

inputs=[
    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
    torch_tensorrt.Input(
        min_shape=(1, 224, 224, 3),
        opt_shape=(1, 512, 512, 3),
        max_shape=(1, 1024, 1024, 3),
        dtype=torch.int32
        format=torch.channel_last
    ), # Dynamic input shape for input #2
    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
]

關鍵字引數

arg_inputs (Tuple[Any, ...]) – 與輸入相同。別名，以便更好地理解 kwarg_inputs。
kwarg_inputs (dict[Any, ...]) – 選擇性，模組 forward 函數的 kwarg 輸入。
device (Union(Device, torch.device, dict)) –
TensorRT 引擎在其上執行的目標裝置
```
device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)
```
disable_tf32 (bool) – 強制 FP32 層使用傳統的 FP32 格式，而不是預設行為（將輸入四捨五入為 10 位元尾數，然後相乘，但使用 23 位元尾數累加總和）
assume_dynamic_shape_support (bool) – 將此設定為 true 可讓轉換器同時適用於動態和靜態形狀。預設值：False
sparse_weights (bool) – 為卷積層和全連接層啟用稀疏性。
enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在選擇核心時可以使用的資料類型集合
debug (bool) – 啟用可偵錯引擎
capability (EngineCapability) – 將核心選擇限制為安全 GPU 核心或安全 DLA 核心
num_avg_timing_iters (python:int) – 用於選擇核心的平均計時迭代次數
workspace_size (python:int) – 提供給 TensorRT 的最大工作區大小
dla_sram_size (python:int) – DLA 用於在層內通訊的快速軟體管理 RAM。
dla_local_dram_size (python:int) – DLA 用於跨操作共用中間張量資料的主機 RAM
dla_global_dram_size (python:int) – DLA 用於儲存權重和執行中繼資料的主機 RAM
truncate_double (bool) – 將以 double (float64) 提供的權重截斷為 float32
calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – 校正器物件，將為 INT8 校正的 PTQ 系統提供資料
require_full_compilation (bool) – 要求模組進行端對端編譯，或傳回錯誤，而不是傳回混合圖形，其中無法在 TensorRT 中執行的操作會在 PyTorch 中執行
min_block_size (python:int) – 為了在 TensorRT 中執行一組操作，連續 TensorRT 可轉換操作的最小數量
torch_executed_ops (Collection[Target]) – 必須在 PyTorch 中執行的 aten 運算子集合。如果此集合不是空的，但 require_full_compilation 為 True，則會擲回錯誤
torch_executed_modules (List[str]) – 必須在 PyTorch 中執行的模組清單。如果此清單不是空的，但 require_full_compilation 為 True，則會擲回錯誤
pass_through_build_failures (bool) – 如果編譯期間發生問題，則傳回錯誤（僅適用於 torch.compile 工作流程）
max_aux_stream (Optional[python:int]) – 引擎中的最大串流數
version_compatible (bool) – 建置與未來 TensorRT 版本相容的 TensorRT 引擎（限制為精簡執行階段運算子，以為引擎提供版本向前相容性）
optimization_level – (Optional[int]): 設定較高的最佳化層級可讓 TensorRT 花費更長的引擎建置時間來搜尋更多最佳化選項。與使用較低最佳化層級建置的引擎相比，產生的引擎可能具有更好的效能。預設最佳化層級為 3。有效值包括從 0 到最大最佳化層級的整數，目前為 5。將其設定為大於最大層級會導致與最大層級相同的行為。
use_python_runtime – (bool): 傳回使用純 Python 執行階段的圖形，減少序列化的選項
use_fast_partitioner – (bool): 使用基於鄰接的分割方案，而不是全域分割器。鄰接分割速度更快，但可能不是最佳的。如果尋求最佳效能，請使用全域分割器 (False)
enable_experimental_decompositions (bool) – 使用完整的運算子分解集合。這些分解可能未經測試，但有助於使圖形更容易轉換為 TensorRT，進而增加在 TensorRT 中執行的圖形數量。
dryrun (bool) – 「Dryrun」模式的切換，執行所有項目，但不包括轉換為 TRT 和記錄輸出
hardware_compatible (bool) – 建置與引擎建置所在 GPU 以外的 GPU 架構相容的 TensorRT 引擎（目前適用於 NVIDIA Ampere 和更新版本）
timing_cache_path (str) – 計時快取的路徑（如果存在），或在編譯後儲存的位置
lazy_engine_init (bool) – 延遲設定引擎，直到所有引擎的編譯完成為止。允許編譯具有多個圖形中斷的較大型模型，但可能導致執行階段 GPU 記憶體過度訂閱。
cache_built_engines (bool) – 是否將編譯的 TRT 引擎儲存到儲存裝置
reuse_cached_engines (bool) – 是否從儲存裝置載入編譯的 TRT 引擎
engine_cache_dir (Optional[str]) – 儲存快取 TRT 引擎的目錄
engine_cache_size (Optional[python:int]) – 用於引擎快取的最大硬碟空間 (位元組)，預設值為 1GB。如果快取超過此大小，預設會移除最舊的引擎
custom_engine_cache (Optional[BaseEngineCache]) – 用於儲存和載入引擎的引擎快取執行個體。使用者可以透過繼承自 BaseEngineCache 來提供自己的引擎快取。如果使用，engine_cache_dir 和 engine_cache_size 將會被忽略。
use_explicit_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) – 啟用權重串流。
**kwargs – 任何項目,

傳回

編譯的 FX 模組，執行時將透過 TensorRT 執行

傳回類型

torch.fx.GraphModule

torch_tensorrt.dynamo.trace(mod: torch.nn.modules.module.Module | torch.fx.graph_module.GraphModule, inputs: Optional[Tuple[Any, ...]] = None, *, arg_inputs: Optional[Tuple[Any, ...]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, **kwargs: Any) → ExportedProgram[source]¶

從 torch.nn.Module 或 torch.fx.GraphModule 匯出 torch.export.ExportedProgram，專門針對使用 Torch-TensorRT 編譯。

從 torch.nn.Module 或 torch.fx.GraphModule 匯出 torch.export.ExportedProgram。執行特定的運算子分解，以利 Torch-TensorRT 的 dynamo 前端進行編譯。

參數

mod (torch.nn.Module | torch.fx.GraphModule) – 來源模組，稍後將由 Torch-TensorRT 的 dynamo 前端編譯

inputs (Tuple[Any, ...]) –

模組輸入的輸入形狀、dtype 和記憶體配置規格清單。此引數為必要項。輸入大小可以指定為 torch 大小、元組或清單。dtype 可以使用 torch 資料類型或 torch_tensorrt 資料類型來指定，您可以使用 torch 裝置或 torch_tensorrt 裝置類型列舉來選擇裝置類型。

input=[
    torch_tensorrt.Input((1, 3, 224, 224)), # Static NCHW input shape for input #1
    torch_tensorrt.Input(
        min_shape=(1, 224, 224, 3),
        opt_shape=(1, 512, 512, 3),
        max_shape=(1, 1024, 1024, 3),
        dtype=torch.int32
        format=torch.channel_last
    ), # Dynamic input shape for input #2
    torch.randn((1, 3, 224, 244)) # Use an example tensor and let torch_tensorrt infer settings
]

關鍵字引數

arg_inputs (Tuple[Any, ...]) – 與輸入相同。別名，以便更好地理解 kwarg_inputs。
kwarg_inputs (dict[Any, ...]) – 選擇性，模組 forward 函數的 kwarg 輸入。
device (Union(torch.device, dict)) –
TensorRT 引擎在其上執行的目標裝置
```
device=torch.device("cuda:0")
```
debug (bool) – 啟用可偵錯引擎
enable_experimental_decompositions (bool) – 使用完整的運算子分解集合。這些分解可能未經測試，但有助於使圖形更容易轉換為 TensorRT，進而增加在 TensorRT 中執行的圖形數量。
**kwargs – 任何項目,

傳回

編譯的 FX 模組，執行時將透過 TensorRT 執行

傳回類型

torch.fx.GraphModule

torch_tensorrt.dynamo.export(gm: GraphModule, cross_compile_flag: Optional[bool] = False) → ExportedProgram[source]¶

將 TensorRT 編譯的結果匯出為所需的輸出格式。

參數

gm (torch.fx.GraphModule) – 已編譯的 Torch-TensorRT 模組，由 torch_tensorrt.dynamo.compile 產生
inputs (torch.Tensor) – Torch 輸入張量
cross_compile_flag (bool) – 指示是否啟用跨編譯的旗標

torch_tensorrt.dynamo.refit_module_weights(compiled_module: torch.fx.graph_module.GraphModule | torch.export.exported_program.ExportedProgram, new_weight_module: ExportedProgram, arg_inputs: Optional[Tuple[Any, ...]] = None, kwarg_inputs: Optional[dict[str, Any]] = None, verify_output: bool = False, use_weight_map_cache: bool = True, in_place: bool = False) → GraphModule[source]¶

使用 ExportedProgram 重新調整已編譯的圖形模組。這會在 compiled_module 中執行權重更新，而無需重新編譯引擎。

參數

compiled_module – 需要重新調整的已編譯 TensorRT 模組。此 compiled_module 應由 torch_tensorrt.dynamo.compile 編譯，或使用 trt.load 從磁碟載入。
new_weight_module – 具有更新權重的匯出程式。此模組應具有與已編譯模組相同的模型架構。
arg_inputs – 範例 arg 輸入。可選，如果需要輸出檢查
kwarg_inputs – 範例 kwarg 輸入。可選，如果需要輸出檢查
verify_output – 是否驗證重新調整模組的輸出

傳回

具有更新權重的新編譯 TensorRT 模組。

類別¶

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, use_aot_joint_export: bool = True)[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) – 為引擎計畫檔案提供版本向前相容性
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 用於跨操作共用中間張量資料的主機 RAM
dla_global_dram_size (python:int) – DLA 用於儲存權重和執行中繼資料的主機 RAM
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 上執行
use_aot_joint_export (bool) – 使用 aot_export_joint_simple，否則使用 AOT_autograd 包裝後端，分散式張量需要此設定

torch_tensorrt.dynamo¶

函數¶

類別¶

文件

教學

資源