ONNX(Open Neural Network Exchange)是由Microsoft、Facebook和AWS共同开发的一种开放格式,旨在促进不同机器学习框架之间的互操作性。ONNX Runtime是一个用于运行转换为ONNX格式的机器学习模型的引擎,无论是传统的机器学习模型还是深度学习模型(神经网络),都可以导出为ONNX格式。ONNX Runtime支持在Linux、Windows和Mac上运行,并且可以在多种芯片架构上运行,包括利用硬件加速器如GPU和TPU。然而,并非所有的操作系统、芯片架构和加速器组合都有现成的安装包,因此,如果使用的不是常见的组合,可能需要从源代码构建运行时。本文将展示如何在x64架构上安装ONNX Runtime,无论是默认CPU还是GPU。
除了能够在多种硬件配置上运行,ONNX Runtime还可以从大多数流行的编程语言调用。本文的目的是展示如何在Python中使用ONNX Runtime。将展示如何安装onnxruntime包。安装ONNX Runtime后,将加载一个之前导出的MNIST模型到ONNX Runtime,并使用它进行预测。
在开始使用ONNX Runtime之前,需要安装onnxruntime包。以下命令将在x64架构上安装带有默认CPU的运行时:
pip install onnxruntime
要在带有GPU的x64架构上安装运行时,请使用以下命令:
pip install onnxruntime-gpu
安装完成后,可以使用以下导入语句将其导入到模块中:
import onnxruntime
将ONNX模型加载到ONNX Runtime中就像转换一样简单,实际上只需要一行代码。下面的函数展示了如何将ONNX模型加载到ONNX Runtime中。实际加载ONNX模型只需要一个命令。
def load_onnx_model(onnx_model_file):
try:
session = onnxruntime.InferenceSession(onnx_model_file)
except (InvalidGraph, TypeError, RuntimeError) as e:
# 可能存在onnxruntime和ONNX模型格式版本之间的不匹配。
print(e)
raise e
return session
下面的函数展示了如何使用加载ONNX模型时创建的ONNX会话。这里有几个值得注意的地方。首先,需要查询会话以获取其输入。这是通过会话的get_inputs()方法完成的。这个输入的名称用于创建一个输入字典,该字典被传递给会话的run命令。MNIST模型只有一个输入参数,即图像列表。(在image_samples参数中。)如果模型有多个输入参数,那么get_inputs()将为每个参数有一个条目。
def onnx_infer(onnx_session, image_samples):
input_name = onnx_session.get_inputs()[0].name
result = onnx_session.run(None, {input_name: image_samples})
probabilities = np.array(result[0])
print(type(probabilities))
print(probabilities)
# 生成预测的arg maxes
predictions = np.argmax(probabilities, axis=1)
return predictions
大多数神经网络不会直接返回预测。它们返回每个输出类别的概率列表。在MNIST模型中,每张图像的返回值将是10个概率的列表。最高概率的条目是预测。一个有趣的测试是将ONNX模型返回的概率与在创建模型的框架内运行原始模型时返回的概率进行比较。理想情况下,模型格式和运行时的更改不应改变产生的任何概率。这将是一个好的单元测试,每次模型发生变化时都会运行。