ONNX Runtime IOBinding 的设备内存复用

我第一次给 GPU 推理做性能拆分时，模型本身并不慢，慢的是每轮 run() 前后的数据搬运：输入从 CPU 复制到设备，输出又复制回 CPU。ONNX Runtime 的 IOBinding 解决的正是这段边界问题。固定 shape 的场景里，把输入和输出都绑定到设备上的 OrtValue，输出缓冲区就能在多轮推理中复用。

官方地址：https://onnxruntime.ai/docs/performance/tune-performance/iobinding.html

IOBinding 先解决的是数据路径，不是模型算子

普通调用把 NumPy 数组作为输入时，ONNX Runtime 需要把数据放到执行提供者使用的设备；如果输出没有预先绑定设备内存，运行结束时又可能把结果搬回 CPU。IOBinding 把这条路径显式化：输入绑定到设备地址，输出绑定到预分配的设备缓冲区，执行阶段只处理图计算。

因此，先用 profiler 或端到端耗时确认瓶颈确实在拷贝。如果模型中仍有 CPU 执行节点，IOBinding 不会把这些节点自动变成 GPU 算子。

固定 shape 时怎样复用设备输出内存

输出维度稳定时，可以提前创建设备上的输出 OrtValue，然后在循环中重复使用。下面示例假定 X_ortvalue 已经由上游设备张量创建，输出形状固定为 [1, 1000]。

import numpy as np
import onnxruntime as ort

# 使用已经配置 CUDAExecutionProvider 的会话，避免每次推理重新创建会话
session = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

# 输入和输出都在 CUDA:0；输出 shape 必须与模型真实输出一致
X_ortvalue = ort.OrtValue.ortvalue_from_numpy(x_cpu, "cuda", 0)
Y_ortvalue = ort.OrtValue.ortvalue_from_shape_and_type(
    [1, 1000], np.float32, "cuda", 0
)

binding = session.io_binding()
for _ in range(10):
    # 输入 shape 不变时，重复绑定同一个输入 OrtValue
    binding.bind_ortvalue_input("input", X_ortvalue)
    # 输出绑定同一个设备缓冲区，推理结果会覆盖上一轮内容
    binding.bind_ortvalue_output("output", Y_ortvalue)
    session.run_with_iobinding(binding)

# 只有需要在 CPU 上处理时才发生一次设备到 CPU 的复制
result = binding.copy_outputs_to_cpu()[0]

这里的“复用”不是把旧结果当成新结果，而是让下一轮把结果写回同一个已分配的缓冲区。若上游框架能原地更新输入，也应保持输入内存和 shape 不变；若输入对象变了，就重新绑定新的 OrtValue。同一个 IOBinding 在复用前要确认绑定状态和输入输出名称没有残留，避免把旧绑定误带入下一轮。

输出 shape 不固定时不要硬塞进旧缓冲区

检测框数量、序列长度或后处理结果会变化时，预分配的 shape 可能不够。此时可只指定设备，让 ONNX Runtime 在设备上分配输出，再通过 get_outputs() 取得设备侧 OrtValue：

binding = session.io_binding()

# 输入仍可绑定已有设备内存；动态输出交给运行时决定 shape
binding.bind_ortvalue_input("input", X_ortvalue)
binding.bind_output("output", "cuda")
session.run_with_iobinding(binding)

# 结果仍在设备上；只有 CPU 后处理需要时才调用 copy_outputs_to_cpu
device_output = binding.get_outputs()[0]
cpu_output = binding.copy_outputs_to_cpu()[0]

动态输出的代价是运行时分配策略可能变化，不能把它等同于固定 buffer 的稳定复用。若下游也支持设备张量，优先消费 device_output；只有日志、NumPy 后处理或接口返回确实需要 CPU 时，才读回。

用这张清单判断是否真的复用了设备内存

排查时把端到端耗时拆成输入准备、run_with_iobinding 和输出读回三段；如果只测模型调用却包含了 copy_outputs_to_cpu()，结论会混入读回成本。

相关问题

IOBinding 会自动消除所有 CPU-GPU 拷贝吗？

不会。它让输入输出的设备位置可控，但 CPU 输入、CPU 后处理、执行提供者不支持的节点仍可能产生拷贝。

为什么输出 buffer 绑定后仍然报 shape 错误？

通常是预分配 shape、dtype 或设备编号与模型输出不一致。先用模型元数据确认输出，再决定固定绑定还是动态分配。

什么时候应调用 copy_outputs_to_cpu？

只有下游逻辑必须拿 NumPy 或 CPU 数据时调用；设备侧后处理可以直接消费 get_outputs() 返回的设备 OrtValue。