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.get_cudagraphs_mode() → bool[source]¶

torch_tensorrt.runtime.get_whole_cudagraphs_mode() → bool[source]¶

torch_tensorrt.runtime.set_cudagraphs_mode(mode: bool) → None[source]¶

torch_tensorrt.runtime.enable_pre_allocated_outputs(module: GraphModule) → _PreAllocatedOutputContextManager[source]¶

torch_tensorrt.runtime.weight_streaming(module: GraphModule) → _WeightStreamingContextManager[source]¶

torch_tensorrt.runtime.enable_output_allocator(module: GraphModule) → _OutputAllocatorContextManager[source]¶

类¶

class torch_tensorrt.runtime.TorchTensorRTModule(**kwargs: Dict[str, Any])[source]¶

TorchTensorRTModule 是一个封装了任意 TensorRT 引擎的 PyTorch 模块。

该模块由 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 引擎实例，负责 [反]序列化、设备配置、剖析
input_binding_names (List[str]) – 输入 TensorRT 引擎绑定名称列表，按传递给 TRT 模块的顺序排列
output_binding_names (List[str]) – 输出 TensorRT 引擎绑定名称列表，按应返回的顺序排列

__init__(**kwargs: Dict[str, Any]) → Any¶: 初始化内部 Module 状态，由 nn.Module 和 ScriptModule 共享。

forward(**kwargs: Dict[str, Any]) → Any¶

定义每次调用时执行的计算。

所有子类都应重写此方法。

注意

尽管 forward 过程的实现需要在该函数内部定义，但之后应该调用 Module 实例而不是它本身，因为前者负责运行已注册的钩子，而后者会静默忽略它们。

get_extra_state(**kwargs: Dict[str, Any]) → Any¶

返回要包含在模块 state_dict 中的任何额外状态。

如果需要存储额外状态，请实现此函数和相应的 set_extra_state()。在构建模块的 state_dict() 时会调用此函数。

注意，为了保证 state_dict 的序列化工作正常，额外状态应该是可被 pickle 的。我们仅为 Tensors 的序列化提供向后兼容性保证；其他对象的序列化形式若发生变化，可能导致向后兼容性中断。

返回: 要存储在模块 state_dict 中的任何额外状态
返回类型: 对象

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>}, 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, offload_module_to_cpu=False), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = None, requires_output_allocator: bool = False, _debugger_config: ~typing.Optional[~torch_tensorrt.dynamo.debug._DebuggerConfig.DebuggerConfig] = None)[source]¶

PythonTorchTensorRTModule 是一个封装了任意 TensorRT 引擎的 PyTorch 模块。

该模块由 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>}, 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, offload_module_to_cpu=False), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = None, requires_output_allocator: bool = False, _debugger_config: ~typing.Optional[~torch_tensorrt.dynamo.debug._DebuggerConfig.DebuggerConfig] = None)[source]¶

接受名称、目标设备、序列化的 TensorRT 引擎以及绑定名称/顺序，并围绕它构建一个 PyTorch torch.nn.Module。使用 TensorRT Python API 运行引擎。

参数

serialized_engine (bytes) – 序列化后的 TensorRT 引擎，形式为字节数组
input_binding_names (List[str]) – 输入 TensorRT 引擎绑定名称列表，按传递给 TRT 模块的顺序排列
output_binding_names (List[str]) – 输出 TensorRT 引擎绑定名称列表，按应返回的顺序排列

关键字参数

name (str) – 模块名称
settings (CompilationSettings) – 用于编译引擎的设置，如果未传入对象，则假定引擎是使用默认编译设置构建的
weight_name_map (dict) – 引擎权重名称到 state_dict 权重名称的映射
requires_output_allocator (bool) – 布尔标志，指示转换器创建的算子是否需要输出分配器来运行（例如，数据依赖性算子）

示例

trt_module = PythonTorchTensorRTModule(
    engine_str,
    input_binding_names=["x"],
    output_binding_names=["output"],
    name="my_module",
    settings=CompilationSettings(device=torch.cuda.current_device)
)

disable_profiling() → None[source]¶: 禁用 TensorRT 剖析。

enable_profiling(profiler: IProfiler = None) → None[source]¶: 启用 TensorRT 剖析。调用此函数后，TensorRT 将在每次前向运行时在标准输出中报告每个层花费的时间。

forward(*inputs: Tensor) → Union[Tensor, Tuple[Tensor, ...]][source]¶

定义每次调用时执行的计算。

所有子类都应重写此方法。

注意

尽管 forward 过程的实现需要在该函数内部定义，但之后应该调用 Module 实例而不是它本身，因为前者负责运行已注册的钩子，而后者会静默忽略它们。

get_layer_info() → str[source]¶: 获取引擎的层信息。仅支持 TRT > 8.2。

validate_input_shapes(inputs: Sequence[Tensor]) → bool[source]¶: 验证前向函数的输入形状是否已更改

torch_tensorrt.runtime¶

函数¶

类¶

文档

教程

资源