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, tiling_optimization_level: str = 'none', l2_limit_for_tiling: int = - 1, **kwargs: Any) → GraphModule[source]¶

使用 TensorRT 为 NVIDIA GPU 编译 ExportedProgram 模块

接受一个现有的 TorchScript 模块和一组配置编译器的设置，并将方法转换为调用等效 TensorRT 引擎的 JIT Graphs

专门转换 TorchScript 模块的 forward 方法

参数

exported_program (torch.export.ExportedProgram) – 源模块，对 torch.nn.Module 运行 torch.export 得到

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

模块输入的输入形状、dtype 和内存布局规格列表。此参数为必需。输入大小可以指定为 torch sizes、tuples 或 lists。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 函数关键字输入。
device (Union(Device, torch.device, dict)) –
TensorRT 引擎的目标运行设备
```
device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)
```
disable_tf32 (布尔值) – 强制 FP32 层使用传统的 FP32 格式，而不是默认的行为（在乘法前将输入舍入到 10 位尾数，但使用 23 位尾数累积求和）
assume_dynamic_shape_support (布尔值) – 将此设置为 true 可使转换器支持动态和静态形状。默认值：False
sparse_weights (布尔值) – 为卷积层和全连接层启用稀疏性。
enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合
debug (布尔值) – 启用可调试引擎
capability (EngineCapability) – 将内核选择限制为安全的 GPU 内核或安全的 DLA 内核
num_avg_timing_iters (整型) – 用于选择内核的平均计时迭代次数
workspace_size (整型) – 提供给 TensorRT 的最大工作空间大小
dla_sram_size (整型) – DLA 用于层内通信的快速软件管理 RAM。
dla_local_dram_size (整型) – DLA 用于跨操作共享中间张量数据的本地 DRAM（主机 RAM）
dla_global_dram_size (整型) – DLA 用于存储权重和元数据以进行执行的全局 DRAM（主机 RAM）
truncate_double (布尔值) – 将以 double (float64) 提供的权重截断为 float32
calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – 校准器对象，将向 PTQ 系统提供数据用于 INT8 校准
require_full_compilation (布尔值) – 要求模块进行端到端编译，否则返回错误，而不是返回一个混合图（其中无法在 TensorRT 中运行的操作在 PyTorch 中运行）
min_block_size (整型) – 为了在 TensorRT 中运行一组操作，连续的可转换为 TensorRT 的操作的最小数量
torch_executed_ops (Collection[Target]) – 必须在 PyTorch 中运行的 aten 运算符集。如果此集合非空但 require_full_compilation 为 True，则会抛出错误
torch_executed_modules (列表[字符串]) – 必须在 PyTorch 中运行的模块列表。如果此列表非空但 require_full_compilation 为 True，则会抛出错误
pass_through_build_failures (布尔值) – 如果编译期间出现问题则报错（仅适用于 torch.compile 工作流程）
max_aux_stream (可选的[整型]) – 引擎中的最大辅助流数量
version_compatible (布尔值) – 构建与未来版本的 TensorRT 兼容的 TensorRT 引擎（限制为精简运行时运算符，以便为引擎提供版本向前兼容性）
optimization_level – (可选的[整型])：设置更高的优化级别允许 TensorRT 花费更长的引擎构建时间来寻找更多优化选项。与较低优化级别构建的引擎相比，由此生成的引擎可能具有更好的性能。默认优化级别为 3。有效值包括从 0 到当前最大优化级别 5 的整数。将其设置为大于最大级别将导致与最大级别相同的行为。
use_python_runtime – (布尔值)：使用纯 Python 运行时返回图，减少了序列化选项
use_fast_partitioner – (布尔值)：使用基于邻接的划分方案而不是全局划分器。邻接划分速度更快，但可能不是最优的。如果寻求最佳性能，请使用全局划分器（False）
enable_experimental_decompositions (布尔值) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易转换为 TensorRT，潜在地增加在 TensorRT 中运行的图的数量。
dryrun (布尔值) – “空运行”模式开关，运行除转换为 TRT 和日志输出之外的所有内容
hardware_compatible (布尔值) – 构建与构建引擎的 GPU 架构之外的其他 GPU 架构兼容的 TensorRT 引擎（目前适用于 NVIDIA Ampere 及更新的架构）
timing_cache_path (字符串) – 如果存在则为计时缓存的路径，或编译后将保存计时缓存的路径
lazy_engine_init (布尔值) – 推迟引擎设置直到所有引擎编译完成。这可以允许具有多个图中断的大型模型进行编译，但在运行时可能导致 GPU 内存过载。
cache_built_engines (布尔值) – 是否将编译好的 TRT 引擎保存到存储中
reuse_cached_engines (布尔值) – 是否从存储加载缓存的 TRT 引擎
engine_cache_dir (可选的[字符串]) – 存储缓存的 TRT 引擎的目录
engine_cache_size (可选的[整型]) – 用于引擎缓存的最大硬盘空间（字节），默认值为 1GB。如果缓存超出此大小，则默认会移除最旧的引擎
custom_engine_cache (Optional[BaseEngineCache]) – 用于保存和加载引擎的引擎缓存实例。用户可以通过继承 BaseEngineCache 提供自己的引擎缓存。如果使用此参数，则 engine_cache_dir 和 engine_cache_size 将被忽略。
use_explicit_typing (布尔值) – 此标志在 TensorRT 编译中启用强类型，该类型会遵守 PyTorch 模型中设置的精度。当用户具有混合精度图时，这很有用。
use_fp32_acc (布尔值) – 此选项在 matmul 层周围插入转换为 FP32 的节点，TensorRT 确保 matmul 的累积在 FP32 中进行。仅当在 enabled_precisions 中配置了 FP16 精度时使用此选项。
refit_identical_engine_weights (布尔值) – 使用相同权重对引擎进行重新适配。当使用不同输入多次编译同一个模型且权重相同时，这很有用。这将通过对不同输入重用相同的引擎来节省时间。
strip_engine_weights (布尔值) – 从序列化的引擎中剥离引擎权重。当引擎部署在不需要权重的环境中时，这很有用。
immutable_weights (布尔值) – 构建不可重新适配的引擎。这对于某些不可重新适配的层很有用。如果此参数设置为 true，则 strip_engine_weights 和 refit_identical_engine_weights 将被忽略。
enable_weight_streaming (布尔值) – 启用权重流式传输。
tiling_optimization_level (字符串) – 平铺策略的优化级别。更高的级别允许 TensorRT 花费更多时间搜索更好的平铺策略。我们当前支持 [“none”, “fast”, “moderate”, “full”]。
l2_limit_for_tiling (整型) – 平铺优化的目标 L2 缓存使用限制（字节）（默认值为 -1，表示无限制）。
**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 sizes、tuples 或 lists。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 函数关键字输入。
device (Union(torch.device, dict)) –
TensorRT 引擎的目标运行设备
```
device=torch.device("cuda:0")
```
debug (布尔值) – 启用可调试引擎
enable_experimental_decompositions (布尔值) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易转换为 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) – PyTorch 输入张量。
cross_compile_flag (bool) – 指示是否启用交叉编译的标志。

torch_tensorrt.dynamo.convert_exported_program_to_serialized_trt_engine(exported_program: ExportedProgram, inputs: Optional[Sequence[Sequence[Any]]] = None, *, arg_inputs: Optional[Sequence[Sequence[Any]]] = None, kwarg_inputs: Optional[dict[Any, Any]] = None, enabled_precisions: Union[Set[Union[dtype, dtype]], Tuple[Union[dtype, dtype]]] = {dtype.f32}, debug: bool = False, assume_dynamic_shape_support: bool = False, workspace_size: int = 0, min_block_size: int = 5, torch_executed_ops: Optional[Set[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: Optional[bool] = False, truncate_double: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, device: Device = Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation: bool = False, disable_tf32: bool = False, sparse_weights: bool = False, engine_capability: EngineCapability = EngineCapability.STANDARD, num_avg_timing_iters: int = 1, dla_sram_size: int = 1048576, dla_local_dram_size: int = 1073741824, dla_global_dram_size: int = 536870912, calibrator: object = None, allow_shape_tensors: bool = False, timing_cache_path: str = '/tmp/torch_tensorrt_engine_cache/timing_cache.bin', 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, tiling_optimization_level: str = 'none', l2_limit_for_tiling: int = - 1, **kwargs: Any) → bytes[source]¶

将 ExportedProgram 转换为序列化的 TensorRT 引擎。

根据转换设置字典，将 ExportedProgram 转换为序列化的 TensorRT 引擎。

参数

exported_program (torch.export.ExportedProgram) – 源模块。

关键字参数

inputs (Optional[Sequence[Input | torch.Tensor]]) –

必需模块输入形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 PyTorch 大小、元组或列表。数据类型可以使用 PyTorch 数据类型或 torch_tensorrt 数据类型指定，您可以使用 PyTorch 设备或 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
  ]

enabled_precisions (Optional[Set[torch.dpython:type | _enums.dpython:type]]) – TensorRT 可以使用的数据类型集合。
debug (bool) – 是否打印详细的调试信息。
workspace_size (python:int) – TensorRT 允许为模块使用的工作空间大小（默认为 0）。
min_block_size (python:int) – 每个 TensorRT 引擎块的最小算子数量。
torch_executed_ops (Set[str]) – 一组无论转换器是否覆盖都将在 PyTorch 中运行的算子。
pass_through_build_failures (bool) – 是否在 TensorRT 引擎构建错误时失败 (True) 或不失败 (False)。
max_aux_streams (Optional[python:int]) – 每个引擎允许的最大辅助 TensorRT 流数量。
version_compatible (bool) – 为引擎计划文件提供版本前向兼容性。
optimization_level (Optional[python:int]) – 构建器优化级别 0-5，级别越高意味着构建时间越长，会搜索更多的优化选项。TensorRT 默认为 3。
use_python_runtime (Optional[bool]) – 是否严格使用 Python 运行时或 C++ 运行时。要根据 C++ 依赖的存在性自动选择运行时（优先选择 C++ 运行时，如果可用），请将此参数留空 None。
truncate_double (bool) – 是否将 float64 TensorRT 引擎输入或权重截断为 float32。
use_fast_partitioner (bool) – 是否使用快速或全局图分区系统。
enable_experimental_decompositions (bool) – 是否启用所有核心 aten 分解，还是只启用选定的子集。
device (Device) – 用于编译模型的 GPU。
require_full_compilation (bool) – 是否要求在 TensorRT 中完全编译图。仅适用于 ir=”dynamo”；对 torch.compile 路径无效。
disable_tf32 (bool) – 是否禁用 TensorRT 层的 TF32 计算。
sparse_weights (bool) – 是否允许构建器使用稀疏权重。
engine_capability (trt.EngineCapability) – 将内核选择限制为安全的 GPU 内核或安全的 DLA 内核。
num_avg_timing_iters (整型) – 用于选择内核的平均计时迭代次数
dla_sram_size (整型) – DLA 用于层内通信的快速软件管理 RAM。
dla_local_dram_size (整型) – DLA 用于跨操作共享中间张量数据的本地 DRAM（主机 RAM）
dla_global_dram_size (整型) – DLA 用于存储权重和元数据以进行执行的全局 DRAM（主机 RAM）
calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – 校准器对象，将向 PTQ 系统提供数据用于 INT8 校准
allow_shape_tensors – （实验性）允许 aten::size 使用 TensorRT 中的 IShapeLayer 输出形状张量。
timing_cache_path (字符串) – 如果存在则为计时缓存的路径，或编译后将保存计时缓存的路径
use_explicit_typing (布尔值) – 此标志在 TensorRT 编译中启用强类型，该类型会遵守 PyTorch 模型中设置的精度。当用户具有混合精度图时，这很有用。
use_fp32_acc (布尔值) – 此选项在 matmul 层周围插入转换为 FP32 的节点，TensorRT 确保 matmul 的累积在 FP32 中进行。仅当在 enabled_precisions 中配置了 FP16 精度时使用此选项。
refit_identical_engine_weights (布尔值) – 使用相同权重对引擎进行重新适配。当使用不同输入多次编译同一个模型且权重相同时，这很有用。这将通过对不同输入重用相同的引擎来节省时间。
strip_engine_weights (布尔值) – 从序列化的引擎中剥离引擎权重。当引擎部署在不需要权重的环境中时，这很有用。
immutable_weights (布尔值) – 构建不可重新适配的引擎。这对于某些不可重新适配的层很有用。如果此参数设置为 true，则 strip_engine_weights 和 refit_identical_engine_weights 将被忽略。
enable_weight_streaming (布尔值) – 启用权重流式传输。
tiling_optimization_level (字符串) – 平铺策略的优化级别。更高的级别允许 TensorRT 花费更多时间搜索更好的平铺策略。我们当前支持 [“none”, “fast”, “moderate”, “full”]。
l2_limit_for_tiling (整型) – 平铺优化的目标 L2 缓存使用限制（字节）（默认值为 -1，表示无限制）。

返回值

序列化的 TensorRT 引擎，可以保存到文件或通过 TensorRT API 反序列化。

返回类型

字节

torch_tensorrt.dynamo.cross_compile_for_windows(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, tiling_optimization_level: str = 'none', l2_limit_for_tiling: int = - 1, **kwargs: Any) → GraphModule[source]¶

在 Linux 中使用 TensorRT 编译 ExportedProgram 模块，以便在 Windows 中进行推理

接受一个导出的程序和一组用于配置编译器的设置，它将把方法转换为调用等效 TensorRT 引擎的 AOT 图

参数

exported_program (torch.export.ExportedProgram) – 源模块，对 torch.nn.Module 运行 torch.export 得到

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

模块输入的输入形状、dtype 和内存布局规格列表。此参数为必需。输入大小可以指定为 torch sizes、tuples 或 lists。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 函数关键字输入。
device (Union(Device, torch.device, dict)) –
TensorRT 引擎的目标运行设备
```
device=torch_tensorrt.Device("dla:1", allow_gpu_fallback=True)
```
disable_tf32 (布尔值) – 强制 FP32 层使用传统的 FP32 格式，而不是默认的行为（在乘法前将输入舍入到 10 位尾数，但使用 23 位尾数累积求和）
assume_dynamic_shape_support (布尔值) – 将此设置为 true 可使转换器支持动态和静态形状。默认值：False
sparse_weights (布尔值) – 为卷积层和全连接层启用稀疏性。
enabled_precision (Set(Union(torch.dpython:type, torch_tensorrt.dpython:type))) – TensorRT 在选择内核时可以使用的数据类型集合
debug (布尔值) – 启用可调试引擎
capability (EngineCapability) – 将内核选择限制为安全的 GPU 内核或安全的 DLA 内核
num_avg_timing_iters (整型) – 用于选择内核的平均计时迭代次数
workspace_size (整型) – 提供给 TensorRT 的最大工作空间大小
dla_sram_size (整型) – DLA 用于层内通信的快速软件管理 RAM。
dla_local_dram_size (整型) – DLA 用于跨操作共享中间张量数据的本地 DRAM（主机 RAM）
dla_global_dram_size (整型) – DLA 用于存储权重和元数据以进行执行的全局 DRAM（主机 RAM）
truncate_double (布尔值) – 将以 double (float64) 提供的权重截断为 float32
calibrator (Union(torch_tensorrt._C.IInt8Calibrator, tensorrt.IInt8Calibrator)) – 校准器对象，将向 PTQ 系统提供数据用于 INT8 校准
require_full_compilation (布尔值) – 要求模块进行端到端编译，否则返回错误，而不是返回一个混合图（其中无法在 TensorRT 中运行的操作在 PyTorch 中运行）
min_block_size (整型) – 为了在 TensorRT 中运行一组操作，连续的可转换为 TensorRT 的操作的最小数量
torch_executed_ops (Collection[Target]) – 必须在 PyTorch 中运行的 aten 运算符集。如果此集合非空但 require_full_compilation 为 True，则会抛出错误
torch_executed_modules (列表[字符串]) – 必须在 PyTorch 中运行的模块列表。如果此列表非空但 require_full_compilation 为 True，则会抛出错误
pass_through_build_failures (布尔值) – 如果编译期间出现问题则报错（仅适用于 torch.compile 工作流程）
max_aux_stream (可选的[整型]) – 引擎中的最大辅助流数量
version_compatible (布尔值) – 构建与未来版本的 TensorRT 兼容的 TensorRT 引擎（限制为精简运行时运算符，以便为引擎提供版本向前兼容性）
optimization_level – (可选的[整型])：设置更高的优化级别允许 TensorRT 花费更长的引擎构建时间来寻找更多优化选项。与较低优化级别构建的引擎相比，由此生成的引擎可能具有更好的性能。默认优化级别为 3。有效值包括从 0 到当前最大优化级别 5 的整数。将其设置为大于最大级别将导致与最大级别相同的行为。
use_python_runtime – (布尔值)：使用纯 Python 运行时返回图，减少了序列化选项
use_fast_partitioner – (布尔值)：使用基于邻接的划分方案而不是全局划分器。邻接划分速度更快，但可能不是最优的。如果寻求最佳性能，请使用全局划分器（False）
enable_experimental_decompositions (布尔值) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易转换为 TensorRT，潜在地增加在 TensorRT 中运行的图的数量。
dryrun (布尔值) – “空运行”模式开关，运行除转换为 TRT 和日志输出之外的所有内容
hardware_compatible (布尔值) – 构建与构建引擎的 GPU 架构之外的其他 GPU 架构兼容的 TensorRT 引擎（目前适用于 NVIDIA Ampere 及更新的架构）
timing_cache_path (字符串) – 如果存在则为计时缓存的路径，或编译后将保存计时缓存的路径
lazy_engine_init (布尔值) – 推迟引擎设置直到所有引擎编译完成。这可以允许具有多个图中断的大型模型进行编译，但在运行时可能导致 GPU 内存过载。
cache_built_engines (布尔值) – 是否将编译好的 TRT 引擎保存到存储中
reuse_cached_engines (布尔值) – 是否从存储加载缓存的 TRT 引擎
engine_cache_dir (可选的[字符串]) – 存储缓存的 TRT 引擎的目录
engine_cache_size (可选的[整型]) – 用于引擎缓存的最大硬盘空间（字节），默认值为 1GB。如果缓存超出此大小，则默认会移除最旧的引擎
custom_engine_cache (Optional[BaseEngineCache]) – 用于保存和加载引擎的引擎缓存实例。用户可以通过继承 BaseEngineCache 提供自己的引擎缓存。如果使用此参数，则 engine_cache_dir 和 engine_cache_size 将被忽略。
use_explicit_typing (布尔值) – 此标志在 TensorRT 编译中启用强类型，该类型会遵守 PyTorch 模型中设置的精度。当用户具有混合精度图时，这很有用。
use_fp32_acc (布尔值) – 此选项在 matmul 层周围插入转换为 FP32 的节点，TensorRT 确保 matmul 的累积在 FP32 中进行。仅当在 enabled_precisions 中配置了 FP16 精度时使用此选项。
refit_identical_engine_weights (布尔值) – 使用相同权重对引擎进行重新适配。当使用不同输入多次编译同一个模型且权重相同时，这很有用。这将通过对不同输入重用相同的引擎来节省时间。
strip_engine_weights (布尔值) – 从序列化的引擎中剥离引擎权重。当引擎部署在不需要权重的环境中时，这很有用。
immutable_weights (布尔值) – 构建不可重新适配的引擎。这对于某些不可重新适配的层很有用。如果此参数设置为 true，则 strip_engine_weights 和 refit_identical_engine_weights 将被忽略。
enable_weight_streaming (布尔值) – 启用权重流式传输。
tiling_optimization_level (字符串) – 平铺策略的优化级别。更高的级别允许 TensorRT 花费更多时间搜索更好的平铺策略。我们当前支持 [“none”, “fast”, “moderate”, “full”]。
l2_limit_for_tiling (整型) – 平铺优化的目标 L2 缓存使用限制（字节）（默认值为 -1，表示无限制）。
**kwargs – 任意类型，

返回值

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

返回类型

torch.fx.GraphModule

torch_tensorrt.dynamo.save_cross_compiled_exported_program(gm: GraphModule, file_path: str) → None[source]¶

将交叉编译的导出程序保存到磁盘。

参数

module (torch.fx.GraphModule) – 交叉编译的 Torch-TensorRT 模块
file_path (str) – 导出程序将保存到磁盘的文件路径

torch_tensorrt.dynamo.load_cross_compiled_exported_program(file_path: str = '') → Any[source]¶

在 Windows 中加载先前在 Linux 中交叉编译的 ExportedProgram 文件

参数: file_path (str) – 磁盘上的文件路径
抛出: ValueError – 如果 API 不是在 Windows 中调用，或者没有文件，或者文件不是有效的 ExportedProgram 文件

torch_tensorrt.dynamo.refit_module_weights(*args: List[Any], **kwargs: Dict[str, Any]) → Any¶

类¶

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, tiling_optimization_level: str = 'none', l2_limit_for_tiling: int = -1, use_distributed_mode_trace: bool = False)[source]¶

Torch-TensorRT Dynamo 路径的编译设置

参数

enabled_precisions (Set[dpython:type]) – 可用的内核 dtype 精度
debug (bool) – 是否打印详细的调试信息。
workspace_size (python:int) – TensorRT 允许为模块使用的工作空间大小（默认为 0）。
min_block_size (python:int) – 每个 TensorRT 引擎块的最小算子数量。
torch_executed_ops (Collection[Target]) – 要在 Torch 中运行的操作集合，无论转换器覆盖率如何
pass_through_build_failures (bool) – 是否在 TensorRT 引擎构建错误时失败 (True) 或不失败 (False)。
max_aux_streams (Optional[python:int]) – 每个引擎允许的最大辅助 TensorRT 流数量。
version_compatible (bool) – 为引擎计划文件提供版本前向兼容性。
optimization_level (Optional[python:int]) – 构建器优化级别 0-5，级别越高意味着构建时间越长，会搜索更多的优化选项。TensorRT 默认为 3。
use_python_runtime (Optional[bool]) – 是否严格使用 Python 运行时或 C++ 运行时。要根据 C++ 依赖的存在性自动选择运行时（优先选择 C++ 运行时，如果可用），请将此参数留空 None。
truncate_double (bool) – 是否将 float64 TensorRT 引擎输入或权重截断为 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 (布尔值) – 将此设置为 true 可使转换器支持动态和静态形状。默认值：False
disable_tf32 (bool) – 是否禁用 TensorRT 层的 TF32 计算。
sparse_weights (bool) – 是否允许构建器使用稀疏权重。
engine_capability (trt.EngineCapability) – 将内核选择限制为安全的 GPU 内核或安全的 DLA 内核。
num_avg_timing_iters (整型) – 用于选择内核的平均计时迭代次数
dla_sram_size (整型) – DLA 用于层内通信的快速软件管理 RAM。
dla_local_dram_size (整型) – DLA 用于跨操作共享中间张量数据的本地 DRAM（主机 RAM）
dla_global_dram_size (整型) – DLA 用于存储权重和元数据以进行执行的全局 DRAM（主机 RAM）
dryrun (Union[bool, str]) – 切换“试运行”模式，该模式会执行所有分区操作，但不进行 TRT 引擎转换。打印图结构和分区性质的详细日志。如果指定了字符串路径，则可选择将输出保存到文件。
hardware_compatible (布尔值) – 构建与构建引擎的 GPU 架构之外的其他 GPU 架构兼容的 TensorRT 引擎（目前适用于 NVIDIA Ampere 及更新的架构）
timing_cache_path (字符串) – 如果存在则为计时缓存的路径，或编译后将保存计时缓存的路径
cache_built_engines (布尔值) – 是否将编译好的 TRT 引擎保存到存储中
reuse_cached_engines (布尔值) – 是否从存储加载缓存的 TRT 引擎
use_strong_typing (bool) – 此标志在 TensorRT 编译中启用强类型检查，它会遵循 Pytorch 模型中设置的精度。这对于用户拥有混合精度图的场景很有用。
use_fp32_acc (布尔值) – 此选项在 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 (布尔值) – 启用权重流式传输。
enable_cross_compile_for_windows (bool) – 默认情况下，此参数为 False，表示 TensorRT 引擎只能在其构建的同一平台上执行。设置为 True 将启用跨平台兼容性，允许引擎在 Linux 上构建并在 Windows 上运行。
tiling_optimization_level (字符串) – 平铺策略的优化级别。更高的级别允许 TensorRT 花费更多时间搜索更好的平铺策略。我们当前支持 [“none”, “fast”, “moderate”, “full”]。
l2_limit_for_tiling (整型) – 平铺优化的目标 L2 缓存使用限制（字节）（默认值为 -1，表示无限制）。
use_distributed_mode_trace (bool) – 使用 aot_autograd 来跟踪图。当分布式模型中存在 DTensors 或分布式张量时启用此功能。

torch_tensorrt.dynamo¶

函数¶

类¶

文档

教程

资源