如何在类似sklearn的onnx(onnx运行时)上进行多次推断

Question

我想在python中使用onnx运行时，根据来自onnx模型的许多输入推断输出。一种方法是使用for循环，但它似乎是一个非常琐碎和缓慢的方法。是否有一种方法可以与学习滑雪的方法相同？

对on运行时的单个预测：

import onnxruntime as ort
sess = ort.InferenceSession("xxxxx.onnx")
input_name = sess.get_inputs()
label_name = sess.get_outputs()[0].name

pred_onnx= sess.run([label_name], {
    input_name[0].name: np.array([[40]]).astype(np.int64),
    input_name[1].name: np.array([[0]]).astype(np.int64),
    input_name[2].name: np.array([[0]]).astype(np.int64)
})
pred_onnx

>> Output: [array([[23]], dtype=float32)]

滑雪板中的单个/多个预测(取决于x_test的大小)：

test_predictions = model.predict(x_test)

Answer 1

最好的方法是ONNX模型支持批处理。根据您提供的输入，它可能已经这样做了。您的3个输入似乎具有形状1,1和您的输出具有形状1,1，这可能意味着第一个维度是批处理大小。带有2,1形状的示例输入看起来类似于[40,50]。

我猜如果你提供两批输入，你就会得到两个输出，所以如下所示

pred_onnx= sess.run([label_name], {
  input_name[0].name: np.array([[40],[40]]).astype(np.int64),
  input_name[1].name: np.array([[0],[0]]).astype(np.int64),
  input_name[2].name: np.array([[0],[0]]).astype(np.int64)
})

可以给出[array([[23],[23]], dtype=float32)]的输出

Answer 2

下面是一个使用批处理推断导出到ONNX的sklearn模型的小型工作示例。

from sklearn import datasets, model_selection, linear_model, pipeline, preprocessing
import numpy as np
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
import onnxruntime
import pandas as pd

# load toy dataset, define sklearn pipeline and fit model
dataset = datasets.load_diabetes()
X, y = dataset.data, dataset.target
X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y)
regr = pipeline.Pipeline(
    [("std", preprocessing.StandardScaler()), ("reg", linear_model.LinearRegression())]
)
regr.fit(X_train, y_train)

# export model to onnx
initial_type = list(
    zip(
        dataset.feature_names,
        [FloatTensorType([None, 1]) for _ in range(len(dataset.feature_names))],
    )
)
onx = convert_sklearn(regr, initial_types=initial_type)
with open("model.onnx", "wb") as f:
    f.write(onx.SerializeToString())

# load model in onnx runtime and make batch inference
df_test = pd.DataFrame(X_test, columns=dataset.feature_names)
sess = onnxruntime.InferenceSession("model.onnx")
inputs = {
    f: df_test[f].astype(np.float32).values.reshape(-1, 1)
    for f in dataset.feature_names
}
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], inputs)[0]


# compare results
regr.predict(X_test)
pred_onx.flatten()

我认为最棘手的部分是得到正确的输入形状进行推理。由于我们指定了FloatTensorType([None, 1])，单个输入数组的形状必须是(x,1)形状，其中x是批数。因此，我们需要将shape (x,)的列值重塑为(x,1)。