如何动态确定ONNX Runtime图像分割模型的GPU最大输入尺寸？

以下是几种安全且有效的方法，帮你确定显卡能承载的最大图像输入尺寸，避免推理时崩溃：

1. 通过形状推理+内存估算提前计算 #

不用实际跑推理，先通过ONNX的形状推理工具推导模型各层张量的形状，再计算总内存占用，结合显卡可用显存估算最大输入尺寸：

步骤1：补全模型张量形状 #

加载ONNX模型后，用形状推理工具补全所有动态张量的形状（需要先指定一个初始输入形状）：

import onnx
from onnxruntime.quantization import shape_inference

# 加载模型
model = onnx.load("my_model.onnx")
# 初始输入形状（格式：[batch_size, channel, height, width]）
input_name = model.graph.input[0].name
initial_input_shape = [1, 3, 256, 256]
# 执行形状推理，得到所有张量的具体形状
inferred_model = shape_inference.infer_shapes(model, {input_name: initial_input_shape})

步骤2：计算模型内存占用 #

遍历模型的输入、输出和中间张量，计算总内存消耗：

def calculate_model_memory(model):
    total_bytes = 0
    # 映射ONNX dtype到字节数
    dtype_byte_map = {
        onnx.TensorProto.FLOAT: 4,
        onnx.TensorProto.FLOAT16: 2,
        onnx.TensorProto.INT32: 4,
        onnx.TensorProto.UINT8: 1
    }

    # 计算中间张量内存
    for tensor in model.graph.value_info:
        dtype = tensor.type.tensor_type.elem_type
        if dtype not in dtype_byte_map:
            continue
        byte_per_elem = dtype_byte_map[dtype]
        # 计算张量总元素数
        elem_count = 1
        for dim in tensor.type.tensor_type.shape.dim:
            elem_count *= dim.dim_value
        total_bytes += elem_count * byte_per_elem

    # 计算输入张量内存
    for input_tensor in model.graph.input:
        dtype = input_tensor.type.tensor_type.elem_type
        if dtype not in dtype_byte_map:
            continue
        byte_per_elem = dtype_byte_map[dtype]
        elem_count = 1
        for dim in input_tensor.type.tensor_type.shape.dim:
            elem_count *= dim.dim_value
        total_bytes += elem_count * byte_per_elem

    return total_bytes / (1024 ** 2)  # 转换为MB

# 获取初始形状下的内存占用
initial_memory_mb = calculate_model_memory(inferred_model)

步骤3：结合显卡显存估算最大尺寸 #

先获取显卡可用显存，再根据输入尺寸与内存的比例关系，计算最大可容纳的图像尺寸：

import torch

# 获取显卡可用显存（MB）
if torch.cuda.is_available():
    free_mem_bytes, _ = torch.cuda.mem_get_info()
    available_mem_mb = free_mem_bytes / (1024 ** 2)
else:
    # CPU环境下调整逻辑，这里默认用户用GPU
    available_mem_mb = 0

# 初始输入的总像素数
initial_pixels = initial_input_shape[2] * initial_input_shape[3]
# 每像素的平均内存占用
mem_per_pixel = initial_memory_mb / initial_pixels

# 预留10%显存作为余量，计算最大可容纳的像素数
max_pixels = (available_mem_mb * 0.9) / mem_per_pixel
# 取正方形尺寸，若模型对尺寸有特殊要求（如必须是32的倍数），可调整为最近的合法值
max_size = int(max_pixels ** 0.5)
# 示例：确保尺寸是32的倍数
max_size = max_size - (max_size % 32)

2. 配置ONNX Runtime显存限制+二分法试错 #

通过SessionOptions设置显存上限，让内存不足时抛出可捕获的异常，再用二分法高效查找最大可行尺寸：

配置显存限制 #

import onnxruntime as ort
import numpy as np

session_options = ort.SessionOptions()
# 设置显存上限（例如8GB，单位为字节）
session_options.set_mem_limit(8 * 1024 * 1024 * 1024)
# 禁用内存模式，减少显存碎片化
session_options.enable_memory_pattern = False

# 创建推理Session
sess = ort.InferenceSession("my_model.onnx", 
                            sess_options=session_options,
                            providers=["CUDAExecutionProvider"])
input_name = sess.get_inputs()[0].name

二分法查找最大尺寸 #

def find_max_input_size(sess, input_name, min_size=256, max_size=2048):
    best_size = min_size
    while min_size <= max_size:
        mid_size = (min_size + max_size) // 2
        # 构造测试输入张量
        test_input = np.random.randn(1, 3, mid_size, mid_size).astype(np.float32)
        try:
            # 尝试推理
            sess.run(None, {input_name: test_input})
            # 成功则尝试更大尺寸
            best_size = mid_size
            min_size = mid_size + 1
        except Exception as e:
            # 失败则尝试更小尺寸
            max_size = mid_size - 1
    return best_size

# 查找最大可行尺寸
max_valid_size = find_max_input_size(sess, input_name)

3. 注意模型输入约束 #

有些模型对输入尺寸有硬性要求：

• 比如下采样/上采样层的步长为2，输入尺寸需为2的幂次；
• 或者模型导出时固定了输入尺寸，需先确认ONNX模型是否支持动态维度（PyTorch导出时需指定dynamic_axes参数）。

内容的提问来源于stack exchange，提问作者rovsmor