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

使用 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, ...]) – 与 inputs 相同。别名，以便更好地理解 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)) – Calibrator 对象，它将为 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 – Any,

返回

编译后的 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

从 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, ...]) – 与 inputs 相同。别名，以便更好地理解 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 – Any,

返回

编译后的 FX 模块，运行时将通过 TensorRT 执行

返回类型

torch.fx.GraphModule

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

将 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

使用 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)

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 包装后端，分布式张量需要