训练、转换和预测模型

训练和部署模型通常涉及以下三个步骤

• 使用 scikit-learn 训练流水线，

• 使用 sklearn-onnx 将其转换为 ONNX，

• 使用 onnxruntime 进行预测。

训练模型

一个使用随机森林和鸢尾花数据集的非常基础的示例。

import skl2onnx
import onnx
import sklearn
from sklearn.linear_model import LogisticRegression
import numpy
import onnxruntime as rt
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx import convert_sklearn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = RandomForestClassifier()
clr.fit(X_train, y_train)
print(clr)

RandomForestClassifier()

将模型转换为 ONNX

initial_type = [("float_input", FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types=initial_type, target_opset=12)

with open("rf_iris.onnx", "wb") as f:
    f.write(onx.SerializeToString())

使用 ONNX Runtime 计算预测结果

sess = rt.InferenceSession("rf_iris.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

[1 2 1 0 1 2 0 0 0 2 1 0 2 1 2 1 1 2 1 1 1 0 0 0 1 1 0 0 2 2 0 0 0 1 0 2 2
 1]

包含逻辑回归的完整示例

clr = LogisticRegression()
clr.fit(X_train, y_train)
initial_type = [("float_input", FloatTensorType([None, X_train.shape[1]]))]
onx = convert_sklearn(clr, initial_types=initial_type, target_opset=12)
with open("logreg_iris.onnx", "wb") as f:
    f.write(onx.SerializeToString())

sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

/home/xadupre/vv/this312/lib/python3.12/site-packages/sklearn/linear_model/_logistic.py:465: ConvergenceWarning: lbfgs failed to converge (status=1):
STOP: TOTAL NO. of ITERATIONS REACHED LIMIT.

Increase the number of iterations (max_iter) or scale the data as shown in:
    https://scikit-learn.cn/stable/modules/preprocessing.html
Please also refer to the documentation for alternative solver options:
    https://scikit-learn.cn/stable/modules/linear_model.html#logistic-regression
  n_iter_i = _check_optimize_result(
[1 2 1 0 1 2 0 0 0 2 1 0 2 1 2 1 1 2 1 1 1 0 0 0 1 1 0 0 2 2 0 0 0 1 0 2 2
 1]

此示例使用的版本

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