使用 Torch-TensorRT 的动态形状

默认情况下，您可以使用不同的输入形状运行 PyTorch 模型，输出形状会即时确定。然而，Torch-TensorRT 是一款 AOT（提前）编译器，它需要关于输入形状的一些预先信息才能编译和优化模型。

使用 torch.export 的动态形状 (AOT)

对于动态输入形状，我们必须提供 (min_shape, opt_shape, max_shape) 参数，以便模型能够针对此输入形状范围进行优化。静态和动态形状的使用示例如下。

注意：以下代码使用 Dynamo 前端。如果使用 Torchscript 前端，请将 ir=dynamo 替换为 ir=ts，行为完全相同。

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
# Compile with static shapes
inputs = torch_tensorrt.Input(shape=[1, 3, 224, 224], dtype=torch.float32)
# or compile with dynamic shapes
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)
trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs)

幕后工作原理

当我们使用 torch_tensorrt.compile API 并指定 ir=dynamo (默认) 时，编译过程分为两个阶段。

• torch_tensorrt.dynamo.trace （使用 torch.export 根据给定的输入跟踪图）

我们使用 torch.export.export() API 来跟踪 PyTorch 模块并将其导出为 torch.export.ExportedProgram。对于动态形状的输入，通过 torch_tensorrt.Input API 提供的 (min_shape, opt_shape, max_shape) 范围用于构造 torch.export.Dim 对象，这些对象在 export API 的 dynamic_shapes 参数中使用。请查看 _tracer.py 文件以了解其幕后工作原理。

• torch_tensorrt.dynamo.compile （使用 TensorRT 编译一个 torch.export.ExportedProgram 对象）

在转换为 TensorRT 时，图已经在节点的元数据中包含了动态形状信息，这些信息将在 engine 构建阶段使用。

自定义动态形状约束

给定输入 x = torch_tensorrt.Input(min_shape, opt_shape, max_shape, dtype)，Torch-TensorRT 尝试在 torch.export 跟踪期间通过相应地使用提供的动态维度构造 torch.export.Dim 对象来自动设置约束。有时，我们可能需要设置额外的约束，如果我们不指定它们，Torchdynamo 会报错。如果您必须为模型设置任何自定义约束（通过使用 torch.export.Dim），我们建议您先导出程序，然后再使用 Torch-TensorRT 进行编译。请参阅此文档，了解如何导出具有动态形状的 PyTorch 模块。这里是一个简单的示例，演示如何导出一个对动态维度施加了一些限制的 matmul 层。

import torch
import torch_tensorrt

class MatMul(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, query, key):
        attn_weight = torch.matmul(query, key.transpose(-1, -2))
        return attn_weight

model = MatMul().eval().cuda()
inputs = [torch.randn(1, 12, 7, 64).cuda(), torch.randn(1, 12, 7, 64).cuda()]
seq_len = torch.export.Dim("seq_len", min=1, max=10)
dynamic_shapes=({2: seq_len}, {2: seq_len})
# Export the model first with custom dynamic shape constraints
exp_program = torch.export.export(model, tuple(inputs), dynamic_shapes=dynamic_shapes)
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs)
# Run inference
trt_gm(*inputs)

使用 torch.compile 的动态形状 (JIT)

torch_tensorrt.compile(model, inputs, ir="torch_compile") 返回一个配置了 TensorRT 后端的 torch.compile 封装函数。在使用 ir=torch_compile 的情况下，用户可以使用 torch._dynamo.mark_dynamic API (https://pytorch.ac.cn/docs/stable/torch.compiler_dynamic_shapes.html) 为输入提供动态形状信息，以避免 TensorRT engine 的重新编译。

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = torch.randn((1, 3, 224, 224), dtype=float32)
# This indicates the dimension 0 is dynamic and the range is [1, 8]
torch._dynamo.mark_dynamic(inputs, 0, min=1, max=8)
trt_gm = torch.compile(model, backend="tensorrt")
# Compilation happens when you call the model
trt_gm(inputs)

# No recompilation of TRT engines with modified batch size
inputs_bs2 = torch.randn((2, 3, 224, 224), dtype=torch.float32)
trt_gm(inputs_bs2)