Note

Click here to download the full example code

Train, convert and predict with ONNX Runtime¶

This example demonstrates an end to end scenario starting with the training of a machine learned model to its use in its converted from.

Train a logistic regression
Conversion to ONNX format
Probabilities
Benchmark with RandomForest

Train a logistic regression ¶

The first step consists in retrieving the iris datset.

from sklearn.datasets import load_iris

iris = load_iris()
X, y = iris.data, iris.target

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y)

Then we fit a model.

from sklearn.linear_model import LogisticRegression

clr = LogisticRegression()
clr.fit(X_train, y_train)

/home/runner/.local/lib/python3.10/site-packages/sklearn/linear_model/_logistic.py:444: ConvergenceWarning: lbfgs failed to converge (status=1):
STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.

Increase the number of iterations (max_iter) or scale the data as shown in:
    https://scikit-learn.org/stable/modules/preprocessing.html
Please also refer to the documentation for alternative solver options:
    https://scikit-learn.org/stable/modules/linear_model.html#logistic-regression
  n_iter_i = _check_optimize_result(

LogisticRegression()

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.

We compute the prediction on the test set and we show the confusion matrix.

from sklearn.metrics import confusion_matrix

pred = clr.predict(X_test)
print(confusion_matrix(y_test, pred))

[[11  0  0]
 [ 0 13  2]
 [ 0  1 11]]

Conversion to ONNX format ¶

We use module sklearn-onnx to convert the model into ONNX format.

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

We load the model with ONNX Runtime and look at its input and output.

import onnxruntime as rt

sess = rt.InferenceSession("logreg_iris.onnx", providers=rt.get_available_providers())

print("input name='{}' and shape={}".format(sess.get_inputs()[0].name, sess.get_inputs()[0].shape))
print("output name='{}' and shape={}".format(sess.get_outputs()[0].name, sess.get_outputs()[0].shape))

input name='float_input' and shape=[None, 4]
output name='output_label' and shape=[None]

We compute the predictions.

input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name

import numpy

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

[[11  0  0]
 [ 0 14  0]
 [ 0  0 13]]

The prediction are perfectly identical.

Probabilities ¶

Probabilities are needed to compute other relevant metrics such as the ROC Curve. Let’s see how to get them first with scikit-learn.

prob_sklearn = clr.predict_proba(X_test)
print(prob_sklearn[:3])

[[1.29030347e-03 5.19977113e-01 4.78732584e-01]
 [7.90572203e-04 3.26327127e-01 6.72882300e-01]
 [9.65026010e-01 3.49738628e-02 1.26697184e-07]]

And then with ONNX Runtime. The probabilies appear to be

prob_name = sess.get_outputs()[1].name
prob_rt = sess.run([prob_name], {input_name: X_test.astype(numpy.float32)})[0]

import pprint

pprint.pprint(prob_rt[0:3])

[{0: 0.0012903024908155203, 1: 0.5199770927429199, 2: 0.4787325859069824},
 {0: 0.0007905727834440768, 1: 0.32632729411125183, 2: 0.6728821396827698},
 {0: 0.9650260210037231, 1: 0.034973856061697006, 2: 1.2669713100876834e-07}]

Let’s benchmark.

from timeit import Timer


def speed(inst, number=10, repeat=20):
    timer = Timer(inst, globals=globals())
    raw = numpy.array(timer.repeat(repeat, number=number))
    ave = raw.sum() / len(raw) / number
    mi, ma = raw.min() / number, raw.max() / number
    print("Average %1.3g min=%1.3g max=%1.3g" % (ave, mi, ma))
    return ave


print("Execution time for clr.predict")
speed("clr.predict(X_test)")

print("Execution time for ONNX Runtime")
speed("sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]")

Execution time for clr.predict
Average 4.4e-05 min=4.25e-05 max=5.47e-05
Execution time for ONNX Runtime
Average 2.23e-05 min=2.16e-05 max=3.08e-05

2.2315799999432783e-05

Let’s benchmark a scenario similar to what a webservice experiences: the model has to do one prediction at a time as opposed to a batch of prediction.

def loop(X_test, fct, n=None):
    nrow = X_test.shape[0]
    if n is None:
        n = nrow
    for i in range(0, n):
        im = i % nrow
        fct(X_test[im : im + 1])


print("Execution time for clr.predict")
speed("loop(X_test, clr.predict, 100)")


def sess_predict(x):
    return sess.run([label_name], {input_name: x.astype(numpy.float32)})[0]


print("Execution time for sess_predict")
speed("loop(X_test, sess_predict, 100)")

Execution time for clr.predict
Average 0.00404 min=0.00402 max=0.00413
Execution time for sess_predict
Average 0.00101 min=0.000995 max=0.00104

0.0010053269700003398

Let’s do the same for the probabilities.

print("Execution time for predict_proba")
speed("loop(X_test, clr.predict_proba, 100)")


def sess_predict_proba(x):
    return sess.run([prob_name], {input_name: x.astype(numpy.float32)})[0]


print("Execution time for sess_predict_proba")
speed("loop(X_test, sess_predict_proba, 100)")

Execution time for predict_proba
Average 0.00613 min=0.0061 max=0.00616
Execution time for sess_predict_proba
Average 0.00106 min=0.00106 max=0.0011

0.0010642362600015076

This second comparison is better as ONNX Runtime, in this experience, computes the label and the probabilities in every case.

Benchmark with RandomForest ¶

We first train and save a model in ONNX format.

from sklearn.ensemble import RandomForestClassifier

rf = RandomForestClassifier()
rf.fit(X_train, y_train)

initial_type = [("float_input", FloatTensorType([1, 4]))]
onx = convert_sklearn(rf, initial_types=initial_type)
with open("rf_iris.onnx", "wb") as f:
    f.write(onx.SerializeToString())

We compare.

sess = rt.InferenceSession("rf_iris.onnx", providers=rt.get_available_providers())

def sess_predict_proba_rf(x):
    return sess.run([prob_name], {input_name: x.astype(numpy.float32)})[0]

print("Execution time for predict_proba")
speed("loop(X_test, rf.predict_proba, 100)")

print("Execution time for sess_predict_proba")
speed("loop(X_test, sess_predict_proba_rf, 100)")

Execution time for predict_proba
Average 0.682 min=0.68 max=0.685
Execution time for sess_predict_proba
Average 0.00132 min=0.0013 max=0.00137

0.001315003579999825

Let’s see with different number of trees.

measures = []

for n_trees in range(5, 51, 5):
    print(n_trees)
    rf = RandomForestClassifier(n_estimators=n_trees)
    rf.fit(X_train, y_train)
    initial_type = [("float_input", FloatTensorType([1, 4]))]
    onx = convert_sklearn(rf, initial_types=initial_type)
    with open("rf_iris_%d.onnx" % n_trees, "wb") as f:
        f.write(onx.SerializeToString())
    sess = rt.InferenceSession("rf_iris_%d.onnx" % n_trees, providers=rt.get_available_providers())

    def sess_predict_proba_loop(x):
        return sess.run([prob_name], {input_name: x.astype(numpy.float32)})[0]

    tsk = speed("loop(X_test, rf.predict_proba, 100)", number=5, repeat=5)
    trt = speed("loop(X_test, sess_predict_proba_loop, 100)", number=5, repeat=5)
    measures.append({"n_trees": n_trees, "sklearn": tsk, "rt": trt})

from pandas import DataFrame

df = DataFrame(measures)
ax = df.plot(x="n_trees", y="sklearn", label="scikit-learn", c="blue", logy=True)
df.plot(x="n_trees", y="rt", label="onnxruntime", ax=ax, c="green", logy=True)
ax.set_xlabel("Number of trees")
ax.set_ylabel("Prediction time (s)")
ax.set_title("Speed comparison between scikit-learn and ONNX Runtime\nFor a random forest on Iris dataset")
ax.legend()

Speed comparison between scikit-learn and ONNX Runtime For a random forest on Iris dataset

5
Average 0.0493 min=0.0492 max=0.0495
Average 0.00102 min=0.001 max=0.00105
10
Average 0.0827 min=0.0826 max=0.0828
Average 0.00103 min=0.00102 max=0.00105
15
Average 0.116 min=0.116 max=0.116
Average 0.00103 min=0.00102 max=0.00105
20
Average 0.149 min=0.149 max=0.149
Average 0.00106 min=0.00105 max=0.00108
25
Average 0.182 min=0.182 max=0.182
Average 0.00107 min=0.00106 max=0.00109
30
Average 0.215 min=0.215 max=0.215
Average 0.00109 min=0.00108 max=0.00112
35
Average 0.249 min=0.248 max=0.249
Average 0.00108 min=0.00107 max=0.00111
40
Average 0.282 min=0.282 max=0.283
Average 0.00111 min=0.0011 max=0.00113
45
Average 0.315 min=0.314 max=0.315
Average 0.00112 min=0.0011 max=0.00114
50
Average 0.348 min=0.348 max=0.348
Average 0.00114 min=0.00113 max=0.00116

<matplotlib.legend.Legend object at 0x7f94760bd870>

Total running time of the script: ( 3 minutes 10.085 seconds)

Gallery generated by Sphinx-Gallery

Train, convert and predict with ONNX Runtime¶

Train a logistic regression ¶

Conversion to ONNX format ¶

Probabilities ¶

Benchmark with RandomForest ¶

ONNX Runtime

Navigation

Related Topics

Train, convert and predict with ONNX Runtime¶

Train a logistic regression¶

Conversion to ONNX format¶

Probabilities¶

Benchmark with RandomForest¶

Train a logistic regression ¶

Conversion to ONNX format ¶

Probabilities ¶

Benchmark with RandomForest ¶