sklearn_logistic_regression_to_onnx_conversion_with_runtime_verification.py

1import numpy as np
2from sklearn.datasets import load_iris
3from sklearn.model_selection import train_test_split
4from sklearn.linear_model import LogisticRegression
5from skl2onnx import convert_sklearn
6from skl2onnx.common.data_types import FloatTensorType
7
8# 1. Load data and train a model
9iris = load_iris()
10X, y = iris.data, iris.target
11X_train, X_test, y_train, y_test = train_test_split(X, y)
12clr = LogisticRegression()
13clr.fit(X_train, y_train)
14
15# 2. Convert the model to ONNX
16initial_type = [('float_input', FloatTensorType([None, 4]))]
17onx = convert_sklearn(clr, initial_types=initial_type)
18
19# 3. Save the model
20with open("logreg_iris.onnx", "wb") as f:
21    f.write(onx.SerializeToString())
22
23# 4. (Optional) Verification using ONNX Runtime
24import onnxruntime as rt
25sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])
26input_name = sess.get_inputs()[0].name
27label_name = sess.get_outputs()[0].name
28pred_onx = sess.run([label_name], {input_name: X_test.astype(np.float32)})[0]
29print(pred_onx)