PyTorch-TensorRT 动态输入形状模型保存指南

动态形状编译与保存的核心要点

在使用PyTorch-TensorRT进行模型加速时，动态输入形状是一个常见需求。开发者经常需要处理不同尺寸的输入数据，而TensorRT提供了灵活的解决方案。

动态形状编译的正确方法

通过torch_tensorrt.Input可以定义输入张量的最小、最优和最大形状范围。例如，对于ResNet50模型，可以这样设置：

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
)]

这种设置允许模型在运行时接受1到8个样本的批量输入，同时保持224x224的图像分辨率。

模型保存的关键细节

当使用torch_tensorrt.save保存编译后的模型时，需要特别注意输入参数的格式要求。与编译时使用Input对象不同，保存时需要提供具体的张量实例。

正确的保存方式应该是：

# 使用opt_shape中定义的尺寸创建输入张量
example_input = torch.randn(4, 3, 224, 224).cuda()
torch_tensorrt.save(trt_gm, "trt_gm.ep", inputs=[example_input])

这个示例输入张量应该在之前定义的min_shape和max_shape范围内，通常使用opt_shape的尺寸最为合适。

技术实现原理

PyTorch-TensorRT在底层会利用提供的示例输入来确定引擎的序列化格式。虽然编译阶段定义了动态范围，但保存操作需要一个具体的张量实例来确保正确的序列化过程。

最佳实践建议

1. 始终验证输入张量尺寸是否在编译时定义的范围内
2. 优先使用opt_shape尺寸作为保存时的输入尺寸
3. 对于多输入模型，确保提供的示例输入列表与模型输入顺序一致
4. 考虑使用torch_tensorrt.load加载模型后，验证不同输入尺寸下的推理结果

通过遵循这些准则，开发者可以充分利用PyTorch-TensorRT的动态形状支持，同时确保模型能够正确保存和加载。