基准测试管道

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

创建管道

我们重用了示例 管道：链接 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 skl2onnx.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.81608913e-19 4.87445968e-10 1.79842280e-08
  3.58700551e-10 1.18138028e-06 4.14411048e-08 1.48275025e-07
  2.50162849e-08 5.51240030e-08]
 [1.37889361e-14 9.99999324e-01 9.17867432e-11 8.30390362e-13
  2.57277808e-07 8.84035057e-12 5.11781442e-11 2.83346412e-11
  4.18965302e-07 1.32796357e-13]]
onnx predict_proba
[[9.99998569e-01 7.81611026e-19 4.87444585e-10 1.79842026e-08
  3.58700042e-10 1.18137802e-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
3.638581280000153
onnxruntime
0.17826826999998957

中间步骤

我们假设最终输出是错误的，我们需要检查管道的每个组件，看看哪个组件出了问题。以下方法修改了 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.9553307900000618
onnxruntime
0.11795157300002757
benchmark <class 'sklearn.linear_model._logistic.LogisticRegression'>
scikit-learn
1.2479175040000428
onnxruntime
0.14069998299987674

此示例使用的版本

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

numpy: 2.3.1
scikit-learn: 1.6.1
onnx:  1.19.0
onnxruntime:  1.23.0
skl2onnx:  1.19.1