torch_tensorrt.runtime¶
函式¶
- torch_tensorrt.runtime.set_multi_device_safe_mode(mode: bool) _MultiDeviceSafeModeContextManager[source]¶
將執行時（僅限 Python，預設）設定為多裝置安全模式
在系統有多個可用裝置的情況下，為了使執行時安全執行，需要進行額外的裝置檢查。這些檢查可能會影響效能，因此是可選加入的。用於抑制在多裝置上下文中不安全執行的警告。
- 引數
mode (bool) – 啟用 (
True) 或停用 (False) 多裝置檢查
示例
with torch_tensorrt.runtime.set_multi_device_safe_mode(True): results = trt_compiled_module(*inputs)
- torch_tensorrt.runtime.enable_cudagraphs(compiled_module: Union[GraphModule, Module]) _CudagraphsContextManager[source]¶
- torch_tensorrt.runtime.enable_pre_allocated_outputs(module: GraphModule) _PreAllocatedOutputContextManager[source]¶
類¶
- class torch_tensorrt.runtime.TorchTensorRTModule(**kwargs: Dict[str, Any])[source]¶
TorchTensorRTModule 是一個 PyTorch 模組，它包含任意 TensorRT Engine。
此模組由 Torch-TensorRT 執行時支援，並且與 FX / Python 部署（只需將
import torch_tensorrt作為應用程式的一部分）以及 TorchScript / C++ 部署完全相容，因為 TorchTensorRTModule 可以傳遞給torch.jit.trace然後儲存。forward 函式非常簡單，為 forward(*args: torch.Tensor) -> Tuple[torch.Tensor]，其中內部實現是
return Tuple(torch.ops.tensorrt.execute_engine(list(inputs), self.engine))> 注意: TorchTensorRTModule 僅支援使用顯式批處理構建的引擎
- 變數
name (str) – 模組名稱（便於除錯）
engine (torch.classes.tensorrt.Engine) – Torch-TensorRT TensorRT Engine 例項，管理[反]序列化、裝置配置、效能分析
input_binding_names (List[str]) – 按傳遞給 TRT 模組的順序排列的輸入 TensorRT 引擎繫結名稱列表
output_binding_names (List[str]) – 按應返回的順序排列的輸出 TensorRT 引擎繫結名稱列表
- __init__(**kwargs: Dict[str, Any]) Any¶
初始化內部 Module 狀態，由 nn.Module 和 ScriptModule 共享。
- forward(**kwargs: Dict[str, Any]) Any¶
定義每次呼叫時執行的計算。
應被所有子類重寫。
注意
雖然前向傳遞的實現需要在該函式中定義，但之後應呼叫
Module例項而不是該函式，因為前者負責執行註冊的鉤子，而後者會靜默忽略它們。
- get_extra_state(**kwargs: Dict[str, Any]) Any¶
返回要包含在模組 state_dict 中的任何額外狀態。
如果您需要儲存額外狀態，請為您自己的模組實現此函式以及相應的
set_extra_state()函式。該函式在構建模組的 state_dict() 時呼叫。請注意，額外狀態應該是可序列化的（picklable），以確保 state_dict 的序列化正常工作。我們只對張量的序列化提供向後相容性保證；如果其他物件的序列化 picklable 形式發生變化，可能會破壞向後相容性。
- 返回
要儲存在模組 state_dict 中的任何額外狀態
- 返回型別
object
- set_extra_state(**kwargs: Dict[str, Any]) Any¶
設定載入的 state_dict 中包含的額外狀態。
此函式從
load_state_dict()呼叫，以處理 state_dict 中找到的任何額外狀態。如果您需要在模組的 state_dict 中儲存額外狀態，請為您自己的模組實現此函式以及相應的get_extra_state()函式。- 引數
state (dict) – 來自 state_dict 的額外狀態
- class torch_tensorrt.runtime.PythonTorchTensorRTModule(serialized_engine: ~typing.Optional[bytes] = None, input_binding_names: ~typing.Optional[~typing.List[str]] = None, output_binding_names: ~typing.Optional[~typing.List[str]] = None, *, name: str = '', settings: ~torch_tensorrt.dynamo._settings.CompilationSettings = CompilationSettings(enabled_precisions={<dtype.f32: 7>}, debug=False, workspace_size=0, min_block_size=5, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_double=False, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, assume_dynamic_shape_support=False, sparse_weights=False, engine_capability=<EngineCapability.STANDARD: 1>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, timing_cache_path='/tmp/torch_tensorrt_engine_cache/timing_cache.bin', lazy_engine_init=False, cache_built_engines=False, reuse_cached_engines=False, use_explicit_typing=False, use_fp32_acc=False, refit_identical_engine_weights=False, strip_engine_weights=False, immutable_weights=True, enable_weight_streaming=False, enable_cross_compile_for_windows=False, tiling_optimization_level='none', l2_limit_for_tiling=-1, use_distributed_mode_trace=False), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = None, requires_output_allocator: bool = False)[source]¶
PythonTorchTensorRTModule 是一個 PyTorch 模組，它包含任意 TensorRT Engine。
此模組由 Torch-TensorRT 執行時支援，僅與 FX / Dynamo / Python 部署相容。此模組無法透過 torch.jit.trace 序列化到 torchscript 以用於 C++ 部署。
- __init__(serialized_engine: ~typing.Optional[bytes] = None, input_binding_names: ~typing.Optional[~typing.List[str]] = None, output_binding_names: ~typing.Optional[~typing.List[str]] = None, *, name: str = '', settings: ~torch_tensorrt.dynamo._settings.CompilationSettings = CompilationSettings(enabled_precisions={<dtype.f32: 7>}, debug=False, workspace_size=0, min_block_size=5, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_double=False, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, assume_dynamic_shape_support=False, sparse_weights=False, engine_capability=<EngineCapability.STANDARD: 1>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, timing_cache_path='/tmp/torch_tensorrt_engine_cache/timing_cache.bin', lazy_engine_init=False, cache_built_engines=False, reuse_cached_engines=False, use_explicit_typing=False, use_fp32_acc=False, refit_identical_engine_weights=False, strip_engine_weights=False, immutable_weights=True, enable_weight_streaming=False, enable_cross_compile_for_windows=False, tiling_optimization_level='none', l2_limit_for_tiling=-1, use_distributed_mode_trace=False), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = None, requires_output_allocator: bool = False)[source]¶
接受名稱、目標裝置、序列化的 TensorRT 引擎以及繫結名稱/順序，並在其周圍構建一個 PyTorch
torch.nn.Module。使用 TensorRT Python API 執行引擎。- 引數
serialized_engine (bytes) – 位元組陣列形式的序列化 TensorRT 引擎
input_binding_names (List[str]) – 按傳遞給 TRT 模組的順序排列的輸入 TensorRT 引擎繫結名稱列表
output_binding_names (List[str]) – 按應返回的順序排列的輸出 TensorRT 引擎繫結名稱列表
- 關鍵字引數
name (str) – 模組名稱
settings (CompilationSettings) – 用於編譯引擎的設定，如果未傳遞物件，則假定引擎是使用預設編譯設定構建的
weight_name_map (dict) – 引擎權重名稱到 state_dict 權重名稱的對映
requires_output_allocator (bool) – 布林標誌，指示轉換器是否建立需要 Output Allocator 才能執行的運算元（例如，資料依賴的運算元）
示例
trt_module = PythonTorchTensorRTModule( engine_str, input_binding_names=["x"], output_binding_names=["output"], name="my_module", settings=CompilationSettings(device=torch.cuda.current_device) )
- enable_profiling(profiler: IProfiler = None) None[source]¶
啟用 TensorRT 效能分析。呼叫此函式後，TensorRT 將在每次前向執行時向標準輸出報告每個層花費的時間。
