训练和部署 scikit-learn 管道

该程序从 scikit-learn 文档中的一个示例开始：绘制单个和投票回归预测，将其转换为 ONNX，最后在不同的运行时计算预测。

训练管道

import numpy
from onnxruntime import InferenceSession
from sklearn.datasets import load_diabetes
from sklearn.ensemble import (
    GradientBoostingRegressor,
    RandomForestRegressor,
    VotingRegressor,
)
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from skl2onnx import to_onnx
from onnx.reference import ReferenceEvaluator


X, y = load_diabetes(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y)

# Train classifiers
reg1 = GradientBoostingRegressor(random_state=1, n_estimators=5)
reg2 = RandomForestRegressor(random_state=1, n_estimators=5)
reg3 = LinearRegression()

ereg = Pipeline(
    steps=[
        ("voting", VotingRegressor([("gb", reg1), ("rf", reg2), ("lr", reg3)])),
    ]
)
ereg.fit(X_train, y_train)

Pipeline(steps=[('voting',
                 VotingRegressor(estimators=[('gb',
                                              GradientBoostingRegressor(n_estimators=5,
                                                                        random_state=1)),
                                             ('rf',
                                              RandomForestRegressor(n_estimators=5,
                                                                    random_state=1)),
                                             ('lr', LinearRegression())]))])

在 Jupyter 环境中，请重新运行此单元格以显示 HTML 表示，或信任笔记本。
在 GitHub 上，HTML 表示无法呈现，请尝试使用 nbviewer.org 加载此页面。

转换模型

第二个参数提供用于训练模型的数据样本。它用于推断 ONNX 图的输入类型。它被转换为单个浮点数，ONNX 运行时可能不支持双精度。

onx = to_onnx(ereg, X_train[:1].astype(numpy.float32), target_opset=12)

使用 ONNX 预测

第一个示例使用 onnxruntime.

sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
pred_ort = sess.run(None, {"X": X_test.astype(numpy.float32)})[0]

pred_skl = ereg.predict(X_test.astype(numpy.float32))

print("Onnx Runtime prediction:\n", pred_ort[:5])
print("Sklearn rediction:\n", pred_skl[:5])

Onnx Runtime prediction:
 [[191.7974 ]
 [245.27979]
 [109.12971]
 [111.75992]
 [173.33968]]
Sklearn rediction:
 [191.79737238 245.27977924 109.12969648 111.75990874 173.33966557]

比较

在部署之前，我们需要比较 scikit-learn 和 ONNX 是否返回相同的预测。

def diff(p1, p2):
    p1 = p1.ravel()
    p2 = p2.ravel()
    d = numpy.abs(p2 - p1)
    return d.max(), (d / numpy.abs(p1)).max()


print(diff(pred_skl, pred_ort))

(2.142160806783977e-05, 1.4564236223492598e-07)

看起来不错。最大的误差（绝对误差和相对误差）在使用浮点数代替双精度引入的误差范围内。我们可以将模型保存为 ONNX 格式，并在使用 onnxruntime 的许多平台上计算相同的预测。

Python 运行时

Python 运行时也可以用于计算预测。它不适合用于生产（它仍然依赖于 python），但对于调查转换失败的原因很有用。

oinf = ReferenceEvaluator(onx)
print(oinf)

ReferenceEvaluator(X) -> variable

它的工作方式几乎相同。

pred_pyrt = oinf.run(None, {"X": X_test.astype(numpy.float32)})[0]
print(diff(pred_skl, pred_pyrt))

(16.51177070735949, 0.11400212547427625)