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

使用TensorRT为NVIDIA GPU编译一个ExportedProgram模块

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

专门转换TorchScript模块的前向方法

Parameters

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

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

List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using torch datatypes or torch_tensorrt datatypes and you can use either torch devices or the torch_tensorrt device type enum to select device type.

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
]

Keyword Arguments

arg_inputs (Tuple[Any, ...]) – 与inputs相同。为了更好地理解kwarg_inputs的别名。
kwarg_inputs (dict[Any, ...]) – 可选的，模块前向函数的kwarg输入。
device (Union(设备, 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) – 启用可调试引擎
能力 (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) – 将双精度（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 (可选[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精度时使用此选项。
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 – 任意,

Returns

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

Return type

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前端进行编译。

Parameters

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)), # 输入 #1 的静态 NCHW 输入形状
    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
    ), # 输入 #2 的动态输入形状
    torch.randn((1, 3, 224, 244)) # 使用示例张量并让 torch_tensorrt 推断设置
]

Keyword Arguments

arg_inputs (Tuple[Any, ...]) – 与inputs相同。为了更好地理解kwarg_inputs的别名。
kwarg_inputs (dict[Any, ...]) – 可选的，模块前向函数的kwarg输入。
device (Union(torch.device, dict)) –
TensorRT引擎运行的目标设备
```
device=torch.device("cuda:0")
```
debug (bool) – 启用可调试引擎
enable_experimental_decompositions (bool) – 使用完整的操作符分解集。这些分解可能未经测试，但有助于使图更容易转换为TensorRT，从而可能增加在TensorRT中运行的图的数量。
**kwargs – 任意,

Returns

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

Return type

torch.fx.GraphModule

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

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

Parameters

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

使用ExportedProgram重新适配已编译的图形模块。这会在不重新编译引擎的情况下在compiled_module中执行权重更新。

Parameters

compiled_module – 需要重新适配的已编译TensorRT模块。这个compiled_module应该由torch_tensorrt.dynamo.compile编译或者使用trt.load从磁盘加载。
new_weight_module – 导出的带有更新权重的程序。这个程序应该具有与编译模块相同的模型架构。
arg_inputs – 示例参数输入。可选，如果需要输出检查则必需
kwarg_inputs – 示例关键字参数输入。可选，如果需要进行输出检查则需要。
verify_output – 是否验证重新拟合模块的输出

Returns

一个新的编译的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)[source]¶

Torch-TensorRT 动态路径的编译设置

Parameters

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 (可选[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 (设备) – 用于编译模型的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) – 此选项在矩阵乘法层周围插入转换为FP32的节点，TensorRT确保矩阵乘法的累加在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上运行。