教程#

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

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

  2. 将模型转换或导出为 ONNX 格式。详情请参阅 ONNX 教程。

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

在本教程中，我们将简要创建一个使用 scikit-learn 的流水线 (pipeline)，将其转换为 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)
print(clr)

>>>

    LogisticRegression()

第 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 numpy
import onnxruntime as rt

sess = rt.InferenceSession(
    "logreg_iris.onnx", providers=rt.get_available_providers())
input_name = sess.get_inputs()[0].name
pred_onx = sess.run(None, {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

>>>

    [1 0 0 2 0 2 0 2 1 2 1 1 0 1 1 0 2 0 2 1 1 2 2 1 1 2 2 0 1 1 2 2 2 0 0 1 2
     1]

通过将特定输出的名称列入清单，可以修改代码以获取该特定的输出。

<<<

import numpy
import onnxruntime as rt

sess = rt.InferenceSession(
    "logreg_iris.onnx", providers=rt.get_available_providers())
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]
print(pred_onx)

>>>

    [1 0 0 2 0 2 0 2 1 2 1 1 0 1 1 0 2 0 2 1 1 2 2 1 1 2 2 0 1 1 2 2 2 0 0 1 2
     1]