torch_tensorrt 如何设置动态的 batch_size?

Question

torch_tensorrt 如何设置动态的 batch_size?

我有一个 pytorch 的 resnet50 网络，我想用 tensorrt+cuda跑

所以我想用 torch_tensorrt 将这个 pytorch 的 resnet50 网络，转成 tensorrt 格式的

但是我的场景是推理，而不是训练，所以输入给模型的 batch_size 不会是一个固定值，可能是 1 也可能是 100等等

import torch_tensorrt
import torch
import torchvision.models as models
import tensorrt as trt
import torch_tensorrt as torch_trt

# 定义 ResNet50 模型
model = models.resnet50(pretrained=True)
model.eval()



batch_size = None
image_channel = 3
image_size = 224

device = torch.device("cuda:0")

model=model.to(device)

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


enabled_precisions = {torch.float}  # Run with fp16

trt_ts_module = torch_tensorrt.compile(
    model,
    inputs=inputs,
    enabled_precisions=enabled_precisions
)

trt_ts_module = trt_ts_module.to(device)


torch.jit.save(
    trt_ts_module, "models/iv_resnet50_export_into_pytorch_tensorrt_model_dynamic.ts")

# pytorch 模型导出成 torch_tensorrt 格式的时候，使用 torch_tensorrt.compile 和 torch.jit.save ，如何指定 batch_size 的大小会动态的

但是我不知道如何转成 tensorRT 格式的时候，设置动态的 batch_size

问了 chatGPT ，他给了我 10 多个方案，都是不发允许的

Answer 1

你想要实现动态 batch size，你要设置一个范围，比如，从 1 到 100。:

inputs = [
    torch_tensorrt.Input(
        min_shape=[1, image_channel, image_size, image_size],
        opt_shape=[1, image_channel, image_size, image_size],
        max_shape=[100, image_channel, image_size, image_size],  # 将最大 batch size 更改为 100
        device=device
    )
]

你根据你的硬件和显存限制，可能要权衡动态 batch size 的范围