ONNX推理性能不及Sklearn？逻辑回归模型转ONNX后耗时过长

我已将Sklearn的Logistic Regression模型对象转换为ONNX模型对象，却发现ONNX的推理（scoring）耗时显著长于sklearn.predict()方法。ONNX号称是优化的推理解决方案，因此我怀疑自己操作有误。且数据集越大，耗时差异越明显，为此我构建了X_large_dataset作为测试样本，执行了如下测试代码：

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import datetime
from sklearn.linear_model import LogisticRegression
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
import numpy as np
import onnxruntime as rt

# create training data
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

# fit model to logistic regression
clr = LogisticRegression()
clr.fit(X_train, y_train)

# convert to onnx format
initial_type = [('float_input', FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types=initial_type)
with open("logreg_iris.onnx", "wb") as f:
    f.write(onx.SerializeToString())
    
# create inference session from onnx object
sess = rt.InferenceSession(
    "logreg_iris.onnx", providers=rt.get_available_providers())
input_name = sess.get_inputs()[0].name

# create a larger dataset as a proxy for large batch processing
X_large_dataset = np.array([[1, 2, 3, 4]]*10_000_000)
start = datetime.datetime.now()
pred_onx = sess.run(None, {input_name: X_large_dataset.astype(np.float32)})[0]
end = datetime.datetime.now()
print("onnx scoring time:", end - start)

# compare to scoring directly with model object
start = datetime.datetime.now()
pred_sk = clr.predict(X_large_dataset)
end = datetime.datetime.now()
print("sklearn scoring time:", end - start)

在我的机器上，这段代码显示Sklearn预测耗时不到1秒，而ONNX耗时达18秒，请问这是什么原因？该如何优化？

原因分析 #

• 数据类型转换与内存拷贝开销：Sklearn的predict默认处理float64数据，测试代码中ONNX推理前将数据转为float32，这一步会产生额外的内存拷贝；同时ONNX Runtime默认情况下在numpy数组与内部张量之间可能存在不必要的内存复制，大批次下这种开销会被放大。
• ONNX Runtime默认配置未优化：默认会话未启用全量图优化，也没有针对CPU/GPU的并行计算进行配置，导致无法充分利用硬件资源。
• 模型特性限制：Logistic Regression是极简单的线性模型，Sklearn的实现基于高度优化的C++底层，而ONNX对于这类轻量模型的优化收益有限，反而会因为框架层的额外开销凸显耗时差异。

优化方案 #

1. 匹配数据类型，减少内存拷贝 #

• 让Sklearn模型与ONNX使用相同的数据类型，初始化模型时指定dtype=np.float32，避免推理时的类型转换：

  clr = LogisticRegression(dtype=np.float32)
  

• 使用rt.OrtValue直接包装numpy数组，跳过内存拷贝：

  X_large_float32 = X_large_dataset.astype(np.float32)
  ort_value = rt.OrtValue.ortvalue_from_numpy(X_large_float32)
  pred_onx = sess.run(None, {input_name: ort_value})[0]
  

2. 优化ONNX Runtime会话配置 #

• 启用全量图优化，并配置多线程加速（根据CPU核心数调整线程数）：

  sess_options = rt.SessionOptions()
  # 启用所有图优化
  sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
  # 设置单算子内部并行线程数
  sess_options.intra_op_num_threads = 8
  # 设置多算子之间并行线程数
  sess_options.inter_op_num_threads = 8
  # 创建会话时传入配置
  sess = rt.InferenceSession("logreg_iris.onnx", sess_options=sess_options, providers=rt.get_available_providers())
  

• 优先使用硬件加速提供者：如果有GPU，指定CUDAExecutionProvider；CPU上确保开启OpenMP相关优化。

3. 优化模型转换过程 #

• 转换时禁用zipmap（避免生成不必要的字典输出，减少处理开销）：

  onx = convert_sklearn(clr, initial_types=initial_type, options={id(clr): {'zipmap': False}})
  

4. 批次拆分处理 #

对于超大批次，将数据拆分为多个小批次进行推理，避免单一批次占用过多内存带宽，同时让ONNX Runtime的并行优化更好发挥作用。

内容的提问来源于stack exchange，提问作者magladde