如何利用torch_tensorrt实现ResNet50模型的动态Batch Size推理？

使用Torch-TensorRT实现ResNet50模型动态Batch Size推理

TensorRT加速PyTorch模型推理时，经常面临Batch Size不固定的挑战。本文详细讲解如何利用Torch-TensorRT库将PyTorch的ResNet50模型转换为支持动态Batch Size推理的TensorRT模型。

问题：开发者希望将预训练ResNet50模型转换为TensorRT模型，但实际应用中输入数据的Batch Size并非固定值（可能从1到100甚至更大）。 原代码使用torch_tensorrt.compile和torch.jit.save导出模型，但无法直接指定动态Batch Size。

解决方案：Torch-TensorRT通过torch_tensorrt.Input类的min_shape、opt_shape和max_shape参数来定义输入张量的形状范围，从而实现动态Batch Size支持。 min_shape代表最小Batch Size，opt_shape代表期望Batch Size，max_shape代表最大Batch Size。 设置这三个参数，告知TensorRT引擎支持的Batch Size范围。

具体实现：原始代码中，inputs参数定义如下：

inputs = [    torch_tensorrt.input(        min_shape=[1, image_channel, image_size, image_size],        opt_shape=[1, image_channel, image_size, image_size],        max_shape=[1, image_channel, image_size, image_size],        device=device    )]

登录后复制