torch_tensorrt 如何设置动态的 batch_size?

作者站长头像
站长
· 阅读数 8

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 多个方案,都是不发允许的

回复
1个回答
avatar
test
2024-07-05

你想要实现动态 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 的范围

回复
likes
适合作为回答的
  • 经过验证的有效解决办法
  • 自己的经验指引,对解决问题有帮助
  • 遵循 Markdown 语法排版,代码语义正确
不该作为回答的
  • 询问内容细节或回复楼层
  • 与题目无关的内容
  • “赞”“顶”“同问”“看手册”“解决了没”等毫无意义的内容