ONNX Runtime IO Binding 怎样减少 GPU 张量拷贝

ONNX Runtime 的 InferenceSession.run() 适合快速验证，但在 CUDA 推理中，如果输入仍是 CPU 上的 NumPy 数组，运行时可能先把输入复制到 GPU；如果输出没有绑定到设备，推理结束后又可能把结果复制回 CPU。减少这两段往返的直接办法是使用 IO Binding：把输入包装成 CUDA 上的 OrtValue，把输出绑定或分配在 CUDA 上，再调用 run_with_iobinding()。

官方地址：https://onnxruntime.ai/

先确认 CUDA 执行提供程序真的在工作

IO Binding 不是替代 CUDA Execution Provider 的开关。创建会话时应把 CUDA 放在 CPU 前面，并先查看输入、输出名称；名称写错时，问题看起来像“绑定没有生效”，实际上是绑定对象与模型节点不匹配。

import onnxruntime as ort

# 中文说明：让 CUDA 优先，CPU 只作为不支持节点的回退执行提供程序。
session = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)
print(session.get_providers())
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

如果环境没有匹配的 CUDA、cuDNN 或 GPU 版运行时，先解决执行提供程序安装和版本兼容；单独增加 io_binding() 不会把 CPU 会话变成 GPU 会话。

让输入和输出停留在 CUDA 设备

最容易理解的路径是先把 NumPy 输入转成 CUDA OrtValue，再绑定到输入。下面的例子把输出也交给 CUDA 分配，适合输出形状运行时才能确定的模型：

import numpy as np
import onnxruntime as ort

# 中文说明：这里的初始转换会把 CPU 数组送到 GPU；生产循环应尽量复用上游 GPU 张量。
x_cpu = np.zeros((1, 3, 224, 224), dtype=np.float32)
x_cuda = ort.OrtValue.ortvalue_from_numpy(x_cpu, "cuda", 0)

io_binding = session.io_binding()
# 中文说明：输入已经有 CUDA 设备内存，绑定时不再走 CPU 输入接口。
io_binding.bind_ortvalue_input(input_name, x_cuda)
# 中文说明：输出形状未知时让运行时在 CUDA 上分配，避免默认回传 CPU。
io_binding.bind_output(output_name, "cuda")
session.run_with_iobinding(io_binding)

# 中文说明：get_outputs 返回设备上的 OrtValue，后续 GPU 后处理可直接消费。
y_cuda = io_binding.get_outputs()[0]

如果输出形状固定，可以预先创建 CUDA OrtValue 并使用 bind_ortvalue_output()。这样更容易复用输出缓冲区；如果形状依赖输入数据，则使用设备绑定并让运行时分配更稳妥。

用表格判断哪一段拷贝仍然存在

因此不要只比较一次调用的耗时。首次创建会话、CUDA 内存分配和模型预热会放大固定成本；更有意义的是在相同 batch、相同输入形状下，分别统计预处理、推理、输出回传和 GPU 后处理。若后续代码最终仍要调用 copy_outputs_to_cpu()，IO Binding 只能消除中间路径的额外往返，不能消除这次最终回传。

常见问题

IO Binding 会不会自动把 NumPy 变成零拷贝 GPU 输入？

不会。CPU NumPy 转成 CUDA OrtValue 仍需要上传；要减少每轮拷贝，应让上游直接产生或复用 GPU 张量。

动态输出为什么更适合 bind_output(name, "cuda")？

因为调用前不必猜测最终形状，运行时可以按实际结果在 CUDA 上分配，再通过 get_outputs() 取得设备上的 OrtValue。

什么时候应该调用 copy_outputs_to_cpu()？

只有当 NumPy、CPU 后处理或外部 CPU 接口确实需要结果时才调用；如果后续仍在 GPU 上计算，就继续传递设备上的输出对象。