sklearn_logistic_regression_iris_to_onnx_conversion.py

1# Train a model.
2from sklearn.datasets import load_iris
3from sklearn.model_selection import train_test_split
4from sklearn.linear_model import LogisticRegression
5iris = load_iris()
6X, y = iris.data, iris.target
7X_train, X_test, y_train, y_test = train_test_split(X, y)
8clr = LogisticRegression()
9clr.fit(X_train, y_train)
10
11# Convert into ONNX format
12from skl2onnx import convert_sklearn
13from skl2onnx.common.data_types import FloatTensorType
14initial_type = [('float_input', FloatTensorType([None, 4]))]
15onx = convert_sklearn(clr, initial_types=initial_type)
16with open("logreg_iris.onnx", "wb") as f:
17    f.write(onx.SerializeToString())