使用 Torch-TensorRT 编译导出程序

Pytorch 2.1 引入了 torch.export API，该 API 可以将 PyTorch 程序中的图导出为 ExportedProgram 对象。Torch-TensorRT dynamo 前端会编译这些 ExportedProgram 对象，并使用 TensorRT 对其进行优化。以下是 dynamo 前端的一个简单用法

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224), dtype=torch.float32).cuda()]
exp_program = torch.export.export(model, tuple(inputs))
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs) # Output is a torch.fx.GraphModule
trt_gm(*inputs)

注意

torch_tensorrt.dynamo.compile 是用户与 Torch-TensorRT dynamo 前端交互的主要 API。模型的输入类型应为 ExportedProgram（理想情况下是 torch.export.export 或 torch_tensorrt.dynamo.trace（将在下一节讨论）的输出），输出类型是 torch.fx.GraphModule 对象。

可配置设置

用户可以通过许多选项自定义使用 TensorRT 进行优化的设置。以下是一些常用的选项：

• inputs - 对于静态形状，它可以是 torch 张量的列表或 torch_tensorrt.Input 对象。对于动态形状，这应该是一个 torch_tensorrt.Input 对象的列表。

• enabled_precisions - TensorRT 构建器在优化期间可以使用的精度集合。

• truncate_long_and_double - 将 long 和 double 值分别截断为 int 和 float。

• torch_executed_ops - 强制由 Torch 执行的操作符。

• min_block_size - 作为 TensorRT 段执行所需的连续操作符的最小数量。

完整的选项列表可以在这里找到

注意

我们目前在 Dynamo 中不支持 INT 精度。目前对 INT 精度的支持存在于我们的 Torchscript IR 中。我们计划在下一个版本中为 dynamo 实现类似的支持。

原理

从底层来看，torch_tensorrt.dynamo.compile 在图上执行以下操作。

• 降级 - 应用降级流程以添加/移除操作符，从而实现最佳转换。

• 分区 - 根据 min_block_size 和 torch_executed_ops 字段将图划分为 PyTorch 和 TensorRT 段。

• 转换 - 在此阶段，PyTorch 操作符被转换为 TensorRT 操作符。

• 优化 - 转换后，我们构建 TensorRT 引擎并将其嵌入到 PyTorch 图中。

追踪

torch_tensorrt.dynamo.trace 可用于追踪 PyTorch 图并生成 ExportedProgram。这在内部执行一些操作符分解，以便进行下游优化。然后可以将 ExportedProgram 与 torch_tensorrt.dynamo.compile API 一起使用。如果您的模型具有动态输入形状，您可以使用此 torch_tensorrt.dynamo.trace 导出具有动态形状的模型。此外，您也可以直接使用带约束的 torch.export。

import torch
import torch_tensorrt

inputs = [torch_tensorrt.Input(min_shape=(1, 3, 224, 224),
                              opt_shape=(4, 3, 224, 224),
                              max_shape=(8, 3, 224, 224),
                              dtype=torch.float32)]
model = MyModel().eval()
exp_program = torch_tensorrt.dynamo.trace(model, inputs)