Train and deploy a scikit-learn pipeline

This program starts from an example in scikit-learn documentation: Plot individual and voting regression predictions, converts it into ONNX and finally computes the predictions a different runtime.

Training a pipeline

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())]))])

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

Converts the model

The second argument gives a sample of the data used to train the model. It is used to infer the input type of the ONNX graph. It is converted into single float and ONNX runtimes may not fully support doubles.

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

Prediction with ONNX

The first example uses 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:
 [[ 95.148315]
 [236.64844 ]
 [192.72498 ]
 [163.4321  ]
 [159.0243  ]]
Sklearn rediction:
 [ 95.14831426 236.64842262 192.7249677  163.4320918  159.02430011]

Comparison

Before deploying, we need to compare that both scikit-learn and ONNX return the same predictions.

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))

(np.float64(2.3916485787367492e-05), np.float64(1.4299860951284095e-07))

It looks good. Biggest errors (absolute and relative) are within the margin error introduced by using floats instead of doubles. We can save the model into ONNX format and compute the same predictions in many platform using onnxruntime.

Python runtime

A python runtime can be used as well to compute the prediction. It is not meant to be used into production (it still relies on python), but it is useful to investigate why the conversion went wrong.

oinf = ReferenceEvaluator(onx)
print(oinf)

ReferenceEvaluator(X) -> variable

It works almost the same way.

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

(np.float64(2.3916485787367492e-05), np.float64(1.4092690429730078e-07))