如何在onnx（onnxruntime）上进行类似于sklearn的多重推断

Question

我想在python中使用onnxruntime根据onnx模型的许多输入推断输出。一种方法是使用for循环，但它似乎是一种非常琐碎且缓慢的方法有没有与sklearn相同的方法？

onnxruntime上的单个预测：

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

sklearn中的单/多预测（取决于x_test的大小）：

test_predictions = model.predict(x_test)

Answer 1

最好的方法是让ONNX模型支持批处理。根据您提供的输入，它可能已经这样做了。您的3个输入似乎具有形状[1,1]，而您的输出具有形状[1,1]，这可能意味着第一个维度是批量大小。形状为[2,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)]