使用 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 对象，该对象用于导出 API 的 dynamic_shapes 参数中。请查看 _tracer.py 文件以了解幕后工作原理。

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

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

自定义动态形状约束

对于给定输入 x = torch_tensorrt.Input(min_shape, opt_shape, max_shape, dtype)，Torch-TensorRT 尝试通过相应地使用提供的动态维度构造 torch.export.Dim 对象，在 torch.export 跟踪期间自动设置约束。有时，我们可能需要设置额外的约束，如果我们没有指定这些约束，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") 返回一个 torch.compile 封装函数，其后端配置为 TensorRT。在 ir=torch_compile 的情况下，用户可以使用 torch._dynamo.mark_dynamic API（https://pytorch.ac.cn/docs/stable/torch.compiler_dynamic_shapes.html）为输入提供动态形状信息，以避免重新编译 TensorRT 引擎。

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)