sklearn to onnx

output.txt

Prediction for [1, 1, 1, 1]: [[0.16937022 0.83062978]]
Prediction for [1, 2, 3, 4]: [[0.25071094 0.74928906]]
The maximum opset needed by this model is only 6.
The maximum opset needed by this model is only 1.
INPUT NAME: float_input
OUTPUTS: ['output_label', 'output_probability']
Prediction with ONNX for [1. 1. 1. 1.]: [[{0: 0.16937017440795898, 1: 0.830629825592041}]]
Prediction with ONNX for [1. 2. 3. 4.]: [[{0: 0.2507109045982361, 1: 0.7492890954017639}]]

sklearn_to_onnx.py

from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification

X, y = make_classification(n_samples=1000, n_features=4,
                           n_informative=2, n_redundant=0,
			               random_state=0, shuffle=False)

clf = RandomForestClassifier(max_depth=2, random_state=0)

clf.fit(X, y)

tests = [ [1, 1, 1, 1], [1, 2, 3, 4] ]
for t in tests:
    result = clf.predict_proba([t])
    print('Prediction for {}: {}'.format(t, result))

# Onnx conversion
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
initial_type = [('float_input', FloatTensorType([1, 4]))]
onx = convert_sklearn(clf, initial_types=initial_type)
with open("iris_clf.onnx", "wb") as f:
    f.write(onx.SerializeToString())

# Run via Onnx Runtime
import numpy as np
import onnxruntime as rt
sess = rt.InferenceSession("iris_clf.onnx")
input_name = sess.get_inputs()[0].name
print('INPUT NAME: {}'.format(input_name))

outputs = sess.get_outputs()
output_names = [o.name for o in sess.get_outputs()]
print('OUTPUTS: {}'.format(output_names))
label_name = output_names[1]  # 'output_probability'

for t in tests:
    X = np.asarray(t).astype(np.float32)
    pred_onnx = sess.run([label_name], {input_name: X})
    print('Prediction with ONNX for {}: {}'.format(X, pred_onnx))