ONNX:可互操作深度学习模型的标准

ONNX:可互操作深度学习模型的标准

了解使用 ONNX 标准跨框架和硬件平台部署模型的好处,在本文中,介绍 ONNX 并通过一个实际示例来解释其巨大潜力。

ONNX 是什么?

ONNX(即开放神经网络交换)是一种用于深度学习模型的开源标准,用来表示深度学习模型的开放格式。所谓开放就是 ONNX 定义了一组与环境、平台均无关的标准格式,来增强各种AI模型的可交互性。是由 Facebook 和 Microsoft 共同开发的,目的是让研究人员和工程师更容易在不同的深度学习框架和硬件平台之间迁移模型。

ONNX 的主要优点之一是它允许轻松地从一个框架(例如 PyTorch)导出模型,并导入到另一个框架(例如 TensorFlow)中。这对于想要尝试不同框架来训练和部署模型的研究人员,或者需要在不同硬件平台上部署模型的工程师特别有吸引力。

ONNX 是什么?

ONNX 还提供了一组用于优化和量化模型的工具,这有助于减少模型的内存和计算需求。这对于在边缘设备和其他资源受限的环境上部署模型特别有用。

ONNX 的另一个重要特点是它得到了广泛的公司和组织的支持。不仅包括 Facebook 和微软,还包括亚马逊、英伟达和英特尔等公司。广泛的支持度确保 ONNX 将继续积极开发和维护,使其成为代表深度学习模型的强大而稳定的标准。

ONNX Runtime

ONNX Runtime 是一个用于执行 ONNX(开放神经网络交换)模型的开源推理引擎。它被设计为高性能和轻量级,非常适合部署在各种硬件平台上,包括边缘设备、服务器和云服务。

ONNX Runtime时提供 C++ APIC# APIPython API 用于执行 ONNX 模型。还提供对多个后端的支持,包括 CUDA 和 OpenCL,这使得它可以在各种硬件平台上运行,例如 NVIDIA GPU 和 Intel CPU。

ONNX Runtime 时非常有用,因为无论要使用什么硬件,都可以通过单个框架使用模型进行推理。

ONNX Runtime时提供 C++ API、C# API 和 Python API 用于执行 ONNX 模型

ONNX Runtime 的主要优势之一是其性能。它使用即时(JIT)编译、内核融合和子图分区等各种技术来优化模型的性能。它还支持分布式部署的线程池和节点间通信,这使其成为大规模部署的合适选择。

ONNX Runtime 还提供对各种模型的支持,包括传统机器学习模型和深度学习模型。使其成为一种多功能推理引擎,可用于从计算机视觉和自然语言处理到语音识别和自动驾驶汽车的广泛应用。

CODE

现在来看一个示例,其中使用经典的 scikit-learn 创建机器学习模型,然后将该模型转换为 ONNX 格式,以便可以将其与 ONNX Runtime 一起使用。

首先,导入必要的库,将模型拉入 sklearn 并导出为经典的 pickle 格式,使用数据集 iris

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
import joblib

# 导入数据
iris = load_iris()
x,y = iris.data, iris.target
x_train, x_test, y_train, y_test = train_test_split(x, y)

# 训练并保存模型 
clr = RandomForestClassifier()
clr.fit(x_train, y_train)
joblib.dump(clr, 'model.pkl', compress = 9)

到这里已经训练并保存了模型,可以重新导入它并将其转换为 ONNX 模型。每个框架都会有自己的转换库,因此,如果在 PyTorch 或 TensorFlow 中开发模型,则必须使用另一种模型。在本例中,使用库 skl2onnx

!pip install skl2onnx
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
import joblib

现在开始转换,不过,需要指定 inital_type,然后创建一个名为 model.onnx 的文件,在其中保存 onnx 模型。

clr = joblib.load('model.pkl')
initial_type = [('float_input', FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types = initial_type)
with open('model.onnx' , 'wb') as f:
  f.write(onx.SerializeToString()) 

现在有了 ONNX 格式的模型,可以导入它并在某些数据上使用它来进行推理。 然后安装 ONNX Runtime。

!pip install onnxruntime
import onnxruntime as rt
import numpy as np

创建数据,并导入模型,从而创建一个会话。指定输入和输出名称(标签),并在数据上运行会话!

data = np.array([[5.4, 6.3, 2.6, 7.4], [3.4, 6.2, 7.4, 2.3],[5.2, 6.4, 4.2,5.6]])

sess = rt.InferenceSession('model.onnx')
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], {input_name: data.astype(np.float32)})[0]
print(pred_onx)

通过利用 ONNX 运行时获得了结果。只需要几个简单的命令!

总结

ONNX 是一种开源标准,可以轻松在不同框架和硬件平台之间移动深度学习模型。它提供了一套用于优化和量化模型的工具,并得到了广泛的公司和组织的支持。因此,ONNX 正在成为深度学习的重要标准,使模型可以轻松共享并跨不同平台部署。