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.8/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))

[[13  0  0]
 [ 0  8  1]
 [ 0  0 16]]

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

[[13  0  0]
 [ 0  8  0]
 [ 0  0 17]]

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

[[3.89179878e-06 1.91737264e-02 9.80822382e-01]
 [1.34575380e-03 4.09726583e-01 5.88927663e-01]
 [1.15008076e-04 5.82656789e-02 9.41619313e-01]]

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: 3.891796040988993e-06, 1: 0.019173720851540565, 2: 0.9808223843574524},
 {0: 0.0013457534369081259, 1: 0.4097265601158142, 2: 0.588927686214447},
 {0: 0.00011500807158881798, 1: 0.058265671133995056, 2: 0.941619336605072}]

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.42e-05 min=4.29e-05 max=5.45e-05
Execution time for ONNX Runtime
Average 1.91e-05 min=1.85e-05 max=2.87e-05

1.9148999999458738e-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.00412 min=0.00409 max=0.00418
Execution time for sess_predict
Average 0.000942 min=0.000935 max=0.000967

0.0009417411449996393

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.00616 min=0.00603 max=0.00695
Execution time for sess_predict_proba
Average 0.000981 min=0.000973 max=0.000989

0.0009807766450003895

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.718 min=0.716 max=0.72
Execution time for sess_predict_proba
Average 0.0011 min=0.00109 max=0.00112

0.0011003508200002443

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.0515 min=0.0512 max=0.0518
Average 0.000957 min=0.000947 max=0.000982
10
Average 0.0866 min=0.0865 max=0.0869
Average 0.000968 min=0.000961 max=0.000987
15
Average 0.122 min=0.121 max=0.122
Average 0.000984 min=0.000977 max=0.000998
20
Average 0.157 min=0.156 max=0.16
Average 0.000992 min=0.000986 max=0.00101
25
Average 0.193 min=0.191 max=0.196
Average 0.00104 min=0.00103 max=0.00105
30
Average 0.226 min=0.226 max=0.226
Average 0.00104 min=0.00103 max=0.00105
35
Average 0.261 min=0.261 max=0.262
Average 0.00106 min=0.00105 max=0.00108
40
Average 0.297 min=0.297 max=0.297
Average 0.00107 min=0.00106 max=0.00109
45
Average 0.333 min=0.331 max=0.336
Average 0.00108 min=0.00107 max=0.00111
50
Average 0.366 min=0.365 max=0.366
Average 0.00108 min=0.00107 max=0.00111

<matplotlib.legend.Legend object at 0x7f832143f460>

Total running time of the script: ( 3 minutes 19.794 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 ¶