ONNX Runtime学习之InferenceSession模块

一、简介

onnxruntime.InferenceSession 是 ONNX Runtime 中用于加载和运行 ONNX 模型的核心模块。它提供了一种灵活的方式来在多种硬件设备（如 CPU、GPU）上执行 ONNX 模型推理。通过 InferenceSession，我们可以高效地加载模型并进行推理操作，适合在生产环境中使用。

二、语法和参数

语法

onnxruntime.InferenceSession(path_or_bytes, sess_options=None, providers=None, provider_options=None)

参数

• path_or_bytes：表示要加载的 ONNX 模型的路径（文件路径）或二进制数据（字节流）。
• sess_options（可选）：SessionOptions 对象，用于配置会话的运行选项。
• providers（可选）：指定模型推理时的硬件提供者列表。例如，['CPUExecutionProvider'] 表示只使用 CPU。
• provider_options（可选）：为每个提供者配置特定的选项。下面列出两个最常用的选项
  • CPUExecutionProvider
  • CUDAExecutionProvider

返回值

返回一个 InferenceSession 对象，代表已加载的 ONNX 模型会话。

三、实例

3.1 加载并使用CPU进行推理

• 代码

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession('model.onnx')

# 创建输入数据
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)

# 获取输入名称
input_name = session.get_inputs()[0].name

# 执行推理
outputs = session.run(None, {input_name: input_data})

# 打印输出结果
print(outputs)

• 输出

[array([[0.1, 0.2, 0.7]], dtype=float32)]

3.2 使用GPU进行推理

• 代码

import onnxruntime as ort
import numpy as np

# 加载ONNX模型并指定GPU提供者
session = ort.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])

# 创建输入数据
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)

# 获取输入名称
input_name = session.get_inputs()[0].name

# 执行推理
outputs = session.run(None, {input_name: input_data})

# 打印输出结果
print(outputs)

• 输出

[array([[0.05, 0.15, 0.8]], dtype=float32)]

四、注意事项

1. InferenceSession 支持多个硬件提供者（如 CPU、GPU），但需要确保指定的提供者在系统中已正确安装并配置。
2. 在执行推理时，输入数据的形状和类型必须与模型的输入要求匹配，否则会抛出错误。
3. 如果模型较大，加载模型可能会消耗较多内存资源，建议根据需要配置硬件提供者。