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)
类¶
- class torch_tensorrt.runtime.TorchTensorRTModule(**kwargs: Dict[str, Any])[source]¶
TorchTensorRTModule 是一个 PyTorch 模块，它包含一个任意的 TensorRT 引擎。
此模块由 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¶
初始化内部模块状态，由 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()` 时调用。
请注意，额外状态应该是可 pickle 的，以确保 state_dict 的正常序列化。我们仅为序列化 Tensor 提供向后兼容性保证；如果其他对象的序列化 pickle 形式发生更改，则可能会破坏向后兼容性。
- 返回
要存储在模块 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>}, 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, use_aot_joint_export=True), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = None)[source]¶
PythonTorchTensorRTModule 是一个 PyTorch 模块，它包含一个任意的 TensorRT 引擎。
此模块由 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, use_aot_joint_export=True), weight_name_map: ~typing.Optional[dict[typing.Any, typing.Any]] = 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 权重名称的映射
示例
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 将在每次前向运行时在 stdout 中报告每层花费的时间。
