Benchmark ONNX conversion#

Example Train and deploy a scikit-learn pipeline converts a simple model. This example takes a similar example but on random data and compares the processing time required by each option to compute predictions.

Training a pipeline#

import numpy
from pandas import DataFrame
from tqdm import tqdm
from sklearn import config_context
from sklearn.datasets import make_regression
from sklearn.ensemble import (
    GradientBoostingRegressor, RandomForestRegressor,
    VotingRegressor)
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from mlprodict.onnxrt import OnnxInference
from onnxruntime import InferenceSession
from skl2onnx import to_onnx
from skl2onnx.tutorial import measure_time


N = 11000
X, y = make_regression(N, n_features=10)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, train_size=0.01)
print("Train shape", X_train.shape)
print("Test shape", X_test.shape)

reg1 = GradientBoostingRegressor(random_state=1)
reg2 = RandomForestRegressor(random_state=1)
reg3 = LinearRegression()
ereg = VotingRegressor([('gb', reg1), ('rf', reg2), ('lr', reg3)])
ereg.fit(X_train, y_train)

Out:

Train shape (110, 10)
Test shape (10890, 10)
VotingRegressor(estimators=[('gb', GradientBoostingRegressor(random_state=1)),
                            ('rf', RandomForestRegressor(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.


Measure the processing time#

We use function skl2onnx.tutorial.measure_time(). The page about assume_finite may be useful if you need to optimize the prediction. We measure the processing time per observation whether or not an observation belongs to a batch or is a single one.

sizes = [(1, 50), (10, 50), (1000, 10), (10000, 5)]

with config_context(assume_finite=True):
    obs = []
    for batch_size, repeat in tqdm(sizes):
        context = {"ereg": ereg, 'X': X_test[:batch_size]}
        mt = measure_time(
            "ereg.predict(X)", context, div_by_number=True,
            number=10, repeat=repeat)
        mt['size'] = context['X'].shape[0]
        mt['mean_obs'] = mt['average'] / mt['size']
        obs.append(mt)

df_skl = DataFrame(obs)
df_skl

Out:

  0%|                                                            | 0/4 [00:00<?, ?it/s]
 25%|#############                                       | 1/4 [00:02<00:06,  2.24s/it]
 50%|##########################                          | 2/4 [00:04<00:04,  2.24s/it]
 75%|#######################################             | 3/4 [00:06<00:02,  2.03s/it]
100%|####################################################| 4/4 [00:11<00:00,  3.16s/it]
100%|####################################################| 4/4 [00:11<00:00,  2.79s/it]
average deviation min_exec max_exec repeat number size mean_obs
0 0.004473 0.000534 0.004291 0.008135 50 10 1 0.004473
1 0.004460 0.000343 0.004332 0.006835 50 10 10 0.000446
2 0.017720 0.001314 0.016942 0.021182 10 10 1000 0.000018
3 0.097764 0.003830 0.090896 0.102486 5 10 10000 0.000010


Graphe.

df_skl.set_index('size')[['mean_obs']].plot(
    title="scikit-learn", logx=True, logy=True)
scikit-learn

ONNX runtime#

The same is done with the two ONNX runtime available.

onx = to_onnx(ereg, X_train[:1].astype(numpy.float32),
              target_opset=14)
sess = InferenceSession(onx.SerializeToString())
oinf = OnnxInference(onx, runtime="python_compiled")

obs = []
for batch_size, repeat in tqdm(sizes):

    # scikit-learn
    context = {"ereg": ereg, 'X': X_test[:batch_size].astype(numpy.float32)}
    mt = measure_time(
        "ereg.predict(X)", context, div_by_number=True,
        number=10, repeat=repeat)
    mt['size'] = context['X'].shape[0]
    mt['skl'] = mt['average'] / mt['size']

    # onnxruntime
    context = {"sess": sess, 'X': X_test[:batch_size].astype(numpy.float32)}
    mt2 = measure_time(
        "sess.run(None, {'X': X})[0]", context, div_by_number=True,
        number=10, repeat=repeat)
    mt['ort'] = mt2['average'] / mt['size']

    # mlprodict
    context = {"oinf": oinf, 'X': X_test[:batch_size].astype(numpy.float32)}
    mt2 = measure_time(
        "oinf.run({'X': X})['variable']", context, div_by_number=True,
        number=10, repeat=repeat)
    mt['pyrt'] = mt2['average'] / mt['size']

    # end
    obs.append(mt)


df = DataFrame(obs)
df

Out:

  0%|                                                            | 0/4 [00:00<?, ?it/s]
 25%|#############                                       | 1/4 [00:02<00:06,  2.32s/it]
 50%|##########################                          | 2/4 [00:05<00:05,  2.56s/it]
 75%|#######################################             | 3/4 [00:07<00:02,  2.35s/it]
100%|####################################################| 4/4 [00:13<00:00,  3.93s/it]
100%|####################################################| 4/4 [00:13<00:00,  3.38s/it]
average deviation min_exec max_exec repeat number size skl ort pyrt
0 0.004481 0.000166 0.004401 0.005481 50 10 1 0.004481 0.000079 0.000067
1 0.005204 0.000348 0.005031 0.007292 50 10 10 0.000520 0.000012 0.000014
2 0.017302 0.000573 0.016954 0.018913 10 10 1000 0.000017 0.000001 0.000002
3 0.095836 0.009330 0.084142 0.111543 5 10 10000 0.000010 0.000001 0.000002


Graph.

df.set_index('size')[['skl', 'ort', 'pyrt']].plot(
    title="Average prediction time per runtime",
    logx=True, logy=True)
Average prediction time per runtime

ONNX runtimes are much faster than scikit-learn to predict one observation. scikit-learn is optimized for training, for batch prediction. That explains why scikit-learn and ONNX runtimes seem to converge for big batches. They use similar implementation, parallelization and languages (C++, openmp).

Total running time of the script: ( 0 minutes 31.317 seconds)

Gallery generated by Sphinx-Gallery