介绍¶

快速入门¶

ONNX Runtime 提供了一种简便的方法，可以在 CPU 或 GPU 上高性能运行机器学习模型，而无需依赖训练框架。机器学习框架通常针对批量训练而非预测进行优化，而预测在应用程序、网站和服务中更为常见。从高层来看，您可以

  1. 使用您喜欢的框架训练模型。

  2. 将模型转换或导出为 ONNX 格式。有关更多详细信息，请参阅 ONNX 教程。

  3. 使用 ONNX Runtime 加载并运行模型。

在本教程中，我们将简要介绍如何使用 scikit-learn 创建流水线，将其转换为 ONNX 格式，并运行首次预测。

步骤 1：使用您喜欢的框架训练模型¶

我们将使用著名的 Iris 数据集。

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

from sklearn.linear_model import LogisticRegression
clr = LogisticRegression()
clr.fit(X_train, y_train)

步骤 2：将模型转换或导出为 ONNX 格式¶

ONNX 是一种描述机器学习模型的格式。它定义了一组常用的运算符来构建模型。有一些工具可以将其他模型格式转换为 ONNX。这里我们将使用 ONNXMLTools。

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

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

步骤 3：使用 ONNX Runtime 加载并运行模型¶

我们将使用 ONNX Runtime 来计算此机器学习模型的预测。

import onnxruntime as rt
sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name

pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]

convert_sklearn、to_onnx、initial_types¶

该模块实现了两个函数： convert_sklearn 和 to_onnx。第一个函数在之前的示例中使用过，它需要两个强制参数

  • 一个 scikit-learn 模型或流水线

  • 初始类型

scikit-learn 不存储有关训练数据集的信息。并非总是能够检索特征数量或它们的类型。这就是为什么该函数需要另一个名为 initial_types 的参数。在许多情况下，训练数据集是一个数值矩阵 X_train。然后它变成 initial_type=[('X', FloatTensorType([None, X_train.shape[1]]))] 。 X 是这个唯一输入的名称，第二个项表示类型和形状。形状是 [None, X_train.shape[1]]，第一个维度是行数，后面是特征数。行数是未定义的，因为在模型转换时请求的预测数量未知。特征数通常是已知的。现在假设输入是一个字符串列后面跟着一个矩阵，那么初始类型将是

initial_type=[
    ('S', StringTensorType([None, 1])),
    ('X', FloatTensorType([None, X_train.shape[1]])),
]

函数 to_onnx 是在与 scikit-learn 核心开发人员讨论后实现的。它还包含一个基于训练数据集的一行来推断正确类型的机制。因此，以下代码 convert_sklearn(clr, initial_types=[('X', FloatTensorType([None, 4]))]) 通常被重写为 to_onnx(clr, X_train[:1]) ，其中 X_train 是训练数据集，它可以是矩阵或数据框。默认情况下输入名称是 'X'，除非 X_train 是数据框。在这种情况下，列名用作输入名称。