ONNX 转换基准测试

示例 训练并部署 scikit-learn 流水线 转换了一个简单的模型。本示例使用了类似的示例，但使用随机数据，并比较了每种选项计算预测所需的处理时间。

训练流水线

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

在 Jupyter 环境中，请重新运行此单元格以显示 HTML 表示或信任该 notebook。
在 GitHub 上，无法渲染 HTML 表示，请尝试使用 nbviewer.org 加载此页面。

测量处理时间

我们使用函数 skl2onnx.tutorial.measure_time()。关于 assume_finite 的页面可能有助于您优化预测。我们测量每个观测值的处理时间，无论该观测值是否属于批次或单个观测值。

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

图表。

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

ONNX runtime

对两种可用的 ONNX runtime 执行相同的操作。

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

图表。

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

ONNX runtime 在预测单个观测值时比 scikit-learn 快得多。scikit-learn 针对训练和批次预测进行了优化。这解释了为什么对于大批量数据，scikit-learn 和 ONNX runtime 的性能似乎趋于一致。它们使用相似的实现、并行化和语言（C++, openmp）。