sklearn_logistic_regression_to_onnx_conversion_with_onnxmltools.py

1import onnxmltools
2from sklearn.datasets import load_iris
3from sklearn.model_selection import train_test_split
4from sklearn.linear_model import LogisticRegression
5from skl2onnx.common.data_types import FloatTensorType
6
7# Load data and train a model
8iris = load_iris()
9X, y = iris.data, iris.target
10X_train, X_test, y_train, y_test = train_test_split(X, y)
11clr = LogisticRegression()
12clr.fit(X_train, y_train)
13
14# Define the input type for the ONNX model
15# The shape [None, 4] means a variable number of rows and 4 columns
16initial_type = [('float_input', FloatTensorType([None, 4]))]
17
18# Convert the scikit-learn model to ONNX
19onnx_model = onnxmltools.convert_sklearn(clr, initial_types=initial_type)
20
21# Save the model
22onnxmltools.utils.save_model(onnx_model, 'logistic_regression.onnx')