TensorRT 后端用于 `torch.compile`¶

本指南介绍了Torch-TensorRT torch.compile 后端：一个深度学习编译器，它使用TensorRT来加速各种模型的JIT风格工作流程。

主要特点¶

Torch-TensorRT torch.compile 后端的主要目标是通过将 torch.compile API 的简单性与 TensorRT 的性能相结合，实现即时编译工作流。调用 torch.compile 后端就像导入 torch_tensorrt 包并指定后端一样简单：

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False)

注意

用户可以使用许多额外的自定义选项。本指南将更深入地讨论这些选项。

后端可以处理各种具有挑战性的模型结构，并提供了一个简单易用的接口，以有效加速模型。此外，它还有许多自定义选项，以确保编译过程适合特定的使用场景。

可自定义设置¶

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

自定义设置使用¶

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
                                options={"truncate_long_and_double": True,
                                         "enabled_precisions": {torch.float, torch.half},
                                         "debug": True,
                                         "min_block_size": 2,
                                         "torch_executed_ops": {"torch.ops.aten.sub.Tensor"},
                                         "optimization_level": 4,
                                         "use_python_runtime": False,})

注意

量化/INT8支持计划在未来的版本中发布；目前，我们支持FP16和FP32精度层。

编译¶

编译是通过将输入传递给模型来触发的，如下所示：

import torch_tensorrt
...
# Causes model compilation to occur
first_outputs = optimized_model(*inputs)

# Subsequent inference runs with the same, or similar inputs will not cause recompilation
# For a full discussion of this, see "Recompilation Conditions" below
second_outputs = optimized_model(*inputs)

编译后¶

编译对象可以在Python会话中用于推理，并将根据下面详细说明的重新编译条件进行重新编译。除了常规推理外，编译过程还可以作为确定模型性能、当前操作符覆盖率和序列化可行性的有用工具。这些点将在下面详细讨论。

模型性能¶

从torch.compile返回的优化模型对于模型基准测试非常有用，因为它可以自动处理编译上下文的变化，或者可能需要重新编译的不同输入。在基准测试不同分布、批量大小或其他标准的输入时，这可以节省时间。

操作符覆盖率¶

编译也是确定特定模型操作符覆盖范围的有用工具。例如，以下编译命令将显示每个图的操作符覆盖范围，但不会编译模型 - 实际上提供了一种“试运行”机制：

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
                                options={"debug": True,
                                         "min_block_size": float("inf"),})

如果您的模型的关键操作符不受支持，请参阅dynamo_conversion以贡献您自己的转换器，或在此处提交问题：https://github.com/pytorch/TensorRT/issues。

序列化的可行性¶

编译也可以帮助展示图的中断和特定模型序列化的可行性。例如，如果一个模型没有图的中断并且使用Torch-TensorRT后端成功编译，那么该模型应该可以通过torch_tensorrt Dynamo IR进行编译和序列化，正如在使用Torch-TensorRT的动态形状中讨论的那样。要确定模型中的图中断数量，torch._dynamo.explain函数非常有用：

import torch
import torch_tensorrt
...
explanation = torch._dynamo.explain(model)(*inputs)
print(f"Graph breaks: {explanation.graph_break_count}")
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False, options={"truncate_long_and_double": True})

动态形状支持¶

Torch-TensorRT torch.compile 后端目前需要为每个遇到的新批次大小重新编译，建议在使用此后端编译时使用 dynamic=False 参数。完整的动态形状支持计划在未来的版本中实现。

重新编译条件¶

一旦模型被编译，后续具有相同形状和数据类型的推理输入，以相同方式遍历图时，将不需要重新编译。此外，每次新的重新编译将在Python会话期间被缓存。例如，如果向模型提供批量大小为4和8的输入，导致两次重新编译，那么在相同会话内的推理过程中，对于这些批量大小的未来输入将不再需要重新编译。计划在未来的版本中支持引擎缓存序列化。

重新编译通常由以下两种事件之一触发：遇到不同大小的输入或以不同方式遍历模型代码的输入。后一种情况可能发生在模型代码包含条件逻辑、复杂循环或数据依赖形状时。torch.compile 在这两种情况下都处理守卫，并确定何时需要重新编译。

TensorRT 后端用于 torch.compile¶