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, make_refittable: bool = False, 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 = 1073741824, custom_engine_cache: Optional[BaseEngineCache] = None, use_explicit_typing: bool = False, use_fp32_acc: 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 (Tuple[Any, ...]) –

模块输入的输入形状、数据类型和内存布局规范列表。此参数是必需的。输入大小可以指定为 torch 大小、元组或列表。数据类型可以使用 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 在选择内核时可以使用的数据类型集
refit (bool) – 启用重新拟合
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)) – 校准器对象，它将为 PTQ 系统提供用于 INT8 校准的数据
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 (可选[str]) – 存储缓存的 TRT 引擎的目录
engine_cache_size (可选[python:int]) – 引擎缓存使用的最大硬盘空间（字节），默认值为 1GB。如果缓存超过此大小，默认情况下将删除最旧的引擎
custom_engine_cache (可选[BaseEngineCache]) – 用于保存和加载引擎的引擎缓存实例。用户可以通过继承 BaseEngineCache 来提供自己的引擎缓存。如果使用，将忽略 engine_cache_dir 和 engine_cache_size。
use_explicit_typing (bool) – 此标志在 TensorRT 编译中启用强类型，它尊重 Pytorch 模型中设置的精度。当用户有混合精度图时，这很有用。
use_fp32_acc (bool) – 此选项在矩阵乘法层周围插入转换为 FP32 的节点，TensorRT 确保矩阵乘法的累积发生在 FP32 中。仅当在 enabled_precisions 中配置 FP16 精度时使用此选项。
**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[source]¶

从 torch.export.ExportedProgram 中导出 torch.nn.Module 或 torch.fx.GraphModule，专门针对使用 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 大小、元组或列表。数据类型可以使用 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, inputs: Sequence[Tensor], kwarg_inputs: Optional[dict[str, Any]] = None) → ExportedProgram[source]¶

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

参数

gm (torch.fx.GraphModule) – 由 torch_tensorrt.dynamo.compile 生成的已编译的 Torch-TensorRT 模块
inputs (torch.Tensor) – Torch 输入张量

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[source]¶

使用 ExportedProgram 重新调整已编译的图模块。这在已编译的模块中执行权重更新，而无需重新编译引擎。

参数

compiled_module – 需要重新调整的已编译的 TensorRT 模块。此编译的模块应由 torch_tensorrt.dynamo.compile 编译或使用 trt.load 从磁盘加载。
new_weight_module – 具有更新权重的已导出程序。它应该具有与已编译模块相同的模型架构。
arg_inputs – 示例参数输入。可选，如果需要输出检查
kwarg_inputs – 示例关键字参数输入。可选，如果需要输出检查
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, make_refittable: 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)[source]¶

Torch-TensorRT Dynamo 路径的编译设置

参数

enabled_precisions (Set[dpython:type]) – 可用的内核数据类型精度
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) – 是否允许构建器使用稀疏权重
refit (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) – 此选项在矩阵乘法层周围插入转换为 FP32 的节点，TensorRT 确保矩阵乘法的累积发生在 FP32 中。仅当在 enabled_precisions 中配置 FP16 精度时使用此选项。

torch_tensorrt.dynamo¶

函数¶

类¶

文档

教程

资源