目录
快捷方式

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, 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, offload_module_to_cpu: bool = False, use_distributed_mode_trace: bool = False, **kwargs: Any) → GraphModule[source]

使用 TensorRT 为 NVIDIA GPU 编译 ExportedProgram 模块

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

特别转换 TorchScript 模块的 forward 方法

参数

  • exported_program (torch.export.ExportedProgram) – 源模块，在 torch.nn.Module 上运行 torch.export

  • inputs (Optional[Sequence[Sequence[Any]]]) –

    模块输入的形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch size、tuple 或 list。数据类型可以使用 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 (Optional[Sequence[Sequence[Any]]]) – 与 inputs 相同。为了更好地与 kwarg_inputs 理解而设置的别名。

  • kwarg_inputs (Optional[dict[Any, Any]]) – 模块前向函数的关键字参数输入。

  • 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_precisions (Union[Set[Union[torch.dpython:type, dpython:type]], Tuple[Union[torch.dpython:type, dpython:type]]]) – TensorRT 在选择内核时可以使用的数据类型集合

  • engine_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

  • require_full_compilation (bool) – 要求模块端到端编译，否则返回错误，而不是返回一个混合图，其中无法在 TensorRT 中运行的操作在 PyTorch 中运行

  • min_block_size (python:int) – 为了在 TensorRT 中运行一组操作，连续的可转换 TensorRT 操作的最小数量

  • torch_executed_ops (Optional[Collection[Target]]) – 必须在 PyTorch 中运行的 aten 运算符集合。如果此集合非空但 require_full_compilation 为 True，则会抛出错误。

  • torch_executed_modules (Optional[List[str]]) – 必须在 PyTorch 中运行的模块列表。如果此列表非空但 require_full_compilation 为 True，则会抛出错误。

  • pass_through_build_failures (bool) – 在编译过程中出现问题时是否报错（仅适用于 torch.compile 工作流）

  • max_aux_streams (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) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易被 Torch-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 (str) – 用于存储缓存的 TRT 引擎的目录

  • engine_cache_size (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) – 启用权重流式传输。

  • tiling_optimization_level (str) – 块状化策略的优化级别。更高的级别允许 TensorRT 花费更多时间来搜索更好的块状化策略。我们目前支持 [“none”, “fast”, “moderate”, “full”]。

  • l2_limit_for_tiling (python:int) – 块状化优化目标 L2 缓存使用限制（以字节为单位）（默认为 -1，表示无限制）。

  • offload_module_to_cpu (bool) – 将模块卸载到 CPU。当需要最小化 GPU 内存使用时，此选项很有用。

  • use_distributed_mode_trace (bool) – 使用 aot_autograd 跟踪图。当分布式模型中存在 DTensors 或分布式张量时，此选项将被启用。

  • **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, ...]) –

    模块输入的形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch size、tuple 或 list。数据类型可以使用 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, ...]) – 可选，模块前向函数的关键字参数输入。

  • device (Union(torch.device, dict)) –

    TensorRT 引擎运行的目标设备

    device=torch.device("cuda:0")

  • debug (bool) – 是否启用可调试引擎

  • enable_experimental_decompositions (bool) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易被 Torch-TensorRT 转换，从而可能增加在 TensorRT 中运行的图的数量。

  • **kwargs – 任意，

返回

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

返回类型

torch.fx.GraphModule

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

将 TensorRT 编译结果导出为所需的输出格式。

参数

  • gm (torch.fx.GraphModule) – 编译后的 Torch-TensorRT 模块，由 torch_tensorrt.dynamo.compile 生成

  • inputs (torch.Tensor) – Torch 输入张量

  • cross_compile_module (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, 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, 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, offload_module_to_cpu: bool = False, use_distributed_mode_trace: bool = False, **kwargs: Any) → bytes[source]

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

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

参数

  • exported_program (torch.export.ExportedProgram) – 源模块，在 torch.nn.Module 上运行 torch.export

  • inputs (Optional[Sequence[Sequence[Any]]]) –

    模块输入的形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch size、tuple 或 list。数据类型可以使用 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 (Optional[Sequence[Sequence[Any]]]) – 与 inputs 相同。为了更好地与 kwarg_inputs 理解而设置的别名。

  • kwarg_inputs (Optional[dict[Any, Any]]) – 模块前向函数的关键字参数输入。

  • 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_precisions (Union[Set[Union[torch.dpython:type, dpython:type]], Tuple[Union[torch.dpython:type, dpython:type]]]) – TensorRT 在选择内核时可以使用的数据类型集合

  • engine_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

  • require_full_compilation (bool) – 要求模块端到端编译，否则返回错误，而不是返回一个混合图，其中无法在 TensorRT 中运行的操作在 PyTorch 中运行

  • min_block_size (python:int) – 为了在 TensorRT 中运行一组操作，连续的可转换 TensorRT 操作的最小数量

  • torch_executed_ops (Optional[Collection[Target]]) – 必须在 PyTorch 中运行的 aten 运算符集合。如果此集合非空但 require_full_compilation 为 True，则会抛出错误。

  • torch_executed_modules (Optional[List[str]]) – 必须在 PyTorch 中运行的模块列表。如果此列表非空但 require_full_compilation 为 True，则会抛出错误。

  • pass_through_build_failures (bool) – 在编译过程中出现问题时是否报错（仅适用于 torch.compile 工作流）

  • max_aux_streams (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) – 使用完整的运算符分解集。这些分解可能未经测试，但有助于使图更容易被 Torch-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 (str) – 用于存储缓存的 TRT 引擎的目录

  • engine_cache_size (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) – 启用权重流式传输。

  • tiling_optimization_level (str) – 块状化策略的优化级别。更高的级别允许 TensorRT 花费更多时间来搜索更好的块状化策略。我们目前支持 [“none”, “fast”, “moderate”, “full”]。

  • l2_limit_for_tiling (python:int) – 块状化优化目标 L2 缓存使用限制（以字节为单位）（默认为 -1，表示无限制）。

  • offload_module_to_cpu (bool) – 将模块卸载到 CPU。当需要最小化 GPU 内存使用时，此选项很有用。

  • use_distributed_mode_trace (bool) – 使用 aot_autograd 跟踪图。当分布式模型中存在 DTensors 或分布式张量时，此选项将被启用。

  • **kwargs – 任意，

返回

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

返回类型

字节

torch_tensorrt.dynamo.cross_compile_for_windows(*args: List[Any], **kwargs: Dict[str, Any]) → Any
torch_tensorrt.dynamo.save_cross_compiled_exported_program(*args: List[Any], **kwargs: Dict[str, Any]) → Any
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>, 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, offload_module_to_cpu: bool = False)[source]

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 上运行。

  • tiling_optimization_level (str) – 块状化策略的优化级别。更高的级别允许 TensorRT 花费更多时间来搜索更好的块状化策略。我们目前支持 [“none”, “fast”, “moderate”, “full”]。

  • l2_limit_for_tiling (python:int) – 块状化优化目标 L2 缓存使用限制（以字节为单位）（默认为 -1，表示无限制）。

  • use_distributed_mode_trace (bool) – 使用 aot_autograd 跟踪图。当分布式模型中存在 DTensors 或分布式张量时，此选项将被启用。

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源

为了分析流量并优化您的体验，我们在此网站上使用 Cookie。点击或浏览即表示您同意我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括可用的控件： Cookie 政策。