Torch-TensorRT Dynamo 后端¶
本指南介绍 Torch-TensorRT Dynamo 后端，它使用 TensorRT 以提前方式优化 Pytorch 模型。
使用 Dynamo 后端¶
Pytorch 2.1 引入了 torch.export 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- 将长整型和双精度型值分别截断为整型和浮点型。torch_executed_ops- 强制由 Torch 执行的运算符。min_block_size- 作为 TensorRT 段执行的连续运算符的最小数量。
可以在 此处 找到所有选项的完整列表
注意
我们目前在 Dynamo 中不支持 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。这在内部执行一些运算符的分解，以供下游优化使用。然后，可以使用 torch_tensorrt.dynamo.compile API 来使用 ExportedProgram。如果您的模型具有动态输入形状，可以使用此 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)
