Torch-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(precision: torch.dtype = torch.float32, debug: bool = False, workspace_size: int = 0, min_block_size: int = 5, torch_executed_ops: typing.Set[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_long_and_double: bool = False, use_fast_partitioner: bool = True, enable_experimental_decompositions: bool = False, device: torch_tensorrt._Device.Device = <factory>, require_full_compilation: bool = False)[source]¶

Torch-TensorRT Dynamo 路径的编译设置

参数

precision (torch.dpython:type) – 模型层的精度
debug (bool) – 是否打印详细的调试信息
workspace_size (python:int) – TRT 允许为模块使用的工作区大小（0 为默认值）
min_block_size (python:int) – 每个 TRT 引擎块的最小操作符数量
torch_executed_ops (Sequence[str]) – 无论转换器覆盖范围如何，都将在 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_long_and_double (bool) – 是否将 int64/float64 TRT 引擎输入或权重截断为 int32/float32
use_fast_partitioner (bool) – 是否使用快速或全局图分区系统
enable_experimental_decompositions (bool) – 是否启用所有核心 aten 分解或仅启用其中一部分
device (Device) – 在其上编译模型的 GPU
require_full_compilation (bool) – 是否要求图在 TensorRT 中完全编译。仅适用于 ir=”dynamo”；对 torch.compile 路径没有影响

自定义设置用法¶

import torch_tensorrt
...
optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False,
                                options={"truncate_long_and_double": True,
                                         "precision": 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 转换器以贡献您自己的转换器，或在此处提交问题： 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 处理这两种情况下的保护，并确定何时需要重新编译。