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 onnx.reference import ReferenceEvaluator
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 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)

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), (100, 10)]

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

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:02<00:04,  2.47s/it]
 67%|██████▋   | 2/3 [00:05<00:02,  2.53s/it]
100%|██████████| 3/3 [00:05<00:00,  1.65s/it]
100%|██████████| 3/3 [00:05<00:00,  1.88s/it]

	average	deviation	min_exec	max_exec	repeat	number	size	mean_obs
0	0.004927	0.000599	0.004307	0.006856	50	10	1	0.004927
1	0.005127	0.000824	0.004372	0.008680	50	10	10	0.000513
2	0.006102	0.000432	0.005597	0.006985	10	10	100	0.000061

Graphe.

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

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(), providers=["CPUExecutionProvider"])
oinf = ReferenceEvaluator(onx)

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

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

    # end
    obs.append(mt)


df = DataFrame(obs)
df

  0%|          | 0/3 [00:00<?, ?it/s]
 33%|███▎      | 1/3 [00:06<00:12,  6.37s/it]
 67%|██████▋   | 2/3 [00:18<00:09,  9.56s/it]
100%|██████████| 3/3 [00:29<00:00, 10.46s/it]
100%|██████████| 3/3 [00:29<00:00,  9.90s/it]

	average	deviation	min_exec	max_exec	repeat	number	size	skl	ort	pyrt
0	0.005081	0.000962	0.004419	0.010415	50	10	1	0.005081	0.000036	0.007617
1	0.004960	0.000507	0.004386	0.007150	50	10	10	0.000496	0.000008	0.001852
2	0.006076	0.000710	0.005279	0.007308	10	10	100	0.000061	0.000004	0.001088

Graph.

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

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