ONNX Runtime学习之InferenceSession模块
ONNXRuntime库学习之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)]
四、注意事项
InferenceSession支持多个硬件提供者（如 CPU、GPU），但需要确保指定的提供者在系统中已正确安装并配置。- 在执行推理时，输入数据的形状和类型必须与模型的输入要求匹配，否则会抛出错误。
- 如果模型较大，加载模型可能会消耗较多内存资源，建议根据需要配置硬件提供者。