基准测试一个管道

以下示例检查管道中的每个步骤，比较并基准测试预测结果。

创建管道

我们重用示例 管道化：链式连接PCA和逻辑回归 中实现的管道。有一个更改是因为 ONNX-ML Imputer 不处理字符串类型。这不能成为最终ONNX管道的一部分，必须被移除。请查找下面以 --- 开头的注释。

import skl2onnx
import onnx
import sklearn
import numpy
from skl2onnx.helpers import collect_intermediate_steps
from timeit import timeit
from skl2onnx.helpers import compare_objects
import onnxruntime as rt
from onnxconverter_common.data_types import FloatTensorType
from skl2onnx import convert_sklearn
import numpy as np
import pandas as pd

from sklearn import datasets
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline

logistic = LogisticRegression()
pca = PCA()
pipe = Pipeline(steps=[("pca", pca), ("logistic", logistic)])

digits = datasets.load_digits()
X_digits = digits.data[:1000]
y_digits = digits.target[:1000]

pipe.fit(X_digits, y_digits)

Pipeline(steps=[('pca', PCA()), ('logistic', LogisticRegression())])

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

转换到ONNX

initial_types = [("input", FloatTensorType((None, X_digits.shape[1])))]
model_onnx = convert_sklearn(pipe, initial_types=initial_types, target_opset=12)

sess = rt.InferenceSession(
    model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]
)
print("skl predict_proba")
print(pipe.predict_proba(X_digits[:2]))
onx_pred = sess.run(None, {"input": X_digits[:2].astype(np.float32)})[1]
df = pd.DataFrame(onx_pred)
print("onnx predict_proba")
print(df.values)

skl predict_proba
[[9.99998530e-01 7.81608915e-19 4.87445983e-10 1.79842282e-08
  3.58700553e-10 1.18138026e-06 4.14411050e-08 1.48275026e-07
  2.50162856e-08 5.51240033e-08]
 [1.37889361e-14 9.99999324e-01 9.17867405e-11 8.30390363e-13
  2.57277806e-07 8.84035067e-12 5.11781433e-11 2.83346409e-11
  4.18965301e-07 1.32796354e-13]]
onnx predict_proba
[[9.99998569e-01 7.81611026e-19 4.87444585e-10 1.79842026e-08
  3.58700042e-10 1.18137689e-06 4.14409520e-08 1.48274751e-07
  2.50162131e-08 5.51239410e-08]
 [1.37888807e-14 9.99999344e-01 9.17865159e-11 8.30387679e-13
  2.57277748e-07 8.84032951e-12 5.11779785e-11 2.83345725e-11
  4.18964021e-07 1.32796280e-13]]

比较输出

compare_objects(pipe.predict_proba(X_digits[:2]), onx_pred)
# No exception so they are the same.

基准测试

print("scikit-learn")
print(timeit("pipe.predict_proba(X_digits[:1])", number=10000, globals=globals()))
print("onnxruntime")
print(
    timeit(
        "sess.run(None, {'input': X_digits[:1].astype(np.float32)})[1]",
        number=10000,
        globals=globals(),
    )
)

scikit-learn
2.1023744829981297
onnxruntime
0.12994074699963676

中间步骤

假设最终输出是错误的，我们需要检查管道的每个组件，看是哪个出了问题。以下方法修改scikit-learn管道以截取中间输出，并为每个运算符生成一个较小的ONNX图。

steps = collect_intermediate_steps(pipe, "pipeline", initial_types)

assert len(steps) == 2

pipe.predict_proba(X_digits[:2])

for _i, step in enumerate(steps):
    onnx_step = step["onnx_step"]
    sess = rt.InferenceSession(
        onnx_step.SerializeToString(), providers=["CPUExecutionProvider"]
    )
    onnx_outputs = sess.run(None, {"input": X_digits[:2].astype(np.float32)})
    skl_outputs = step["model"]._debug.outputs
    if "transform" in skl_outputs:
        compare_objects(skl_outputs["transform"], onnx_outputs[0])
        print("benchmark", step["model"].__class__)
        print("scikit-learn")
        print(
            timeit(
                "step['model'].transform(X_digits[:1])", number=10000, globals=globals()
            )
        )
    else:
        compare_objects(skl_outputs["predict_proba"], onnx_outputs[1])
        print("benchmark", step["model"].__class__)
        print("scikit-learn")
        print(
            timeit(
                "step['model'].predict_proba(X_digits[:1])",
                number=10000,
                globals=globals(),
            )
        )
    print("onnxruntime")
    print(
        timeit(
            "sess.run(None, {'input': X_digits[:1].astype(np.float32)})",
            number=10000,
            globals=globals(),
        )
    )

benchmark <class 'sklearn.decomposition._pca.PCA'>
scikit-learn
0.6131834360021458
onnxruntime
0.07247600700065959
benchmark <class 'sklearn.linear_model._logistic.LogisticRegression'>
scikit-learn
0.8395350790015073
onnxruntime
0.10347538899804931

此示例使用的版本

print("numpy:", numpy.__version__)
print("scikit-learn:", sklearn.__version__)
print("onnx: ", onnx.__version__)
print("onnxruntime: ", rt.__version__)
print("skl2onnx: ", skl2onnx.__version__)

numpy: 2.2.0
scikit-learn: 1.6.0
onnx:  1.18.0
onnxruntime:  1.21.0+cu126
skl2onnx:  1.18.0