Choose appropriate output of a classifier

A scikit-learn classifier usually returns a matrix of probabilities. By default, sklearn-onnx converts that matrix into a list of dictionaries where each probabily is mapped to its class id or name. That mechanism retains the class names but is slower. Let’s see what other options are available.

Train a model and convert it

from timeit import repeat
import numpy
import sklearn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import onnxruntime as rt
import onnx
import skl2onnx
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx import to_onnx
from sklearn.linear_model import LogisticRegression
from sklearn.multioutput import MultiOutputClassifier

iris = load_iris()
X, y = iris.data, iris.target
X = X.astype(numpy.float32)
y = y * 2 + 10  # to get labels different from [0, 1, 2]
X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = LogisticRegression(max_iter=500)
clr.fit(X_train, y_train)
print(clr)

onx = to_onnx(clr, X_train, target_opset=12)

LogisticRegression(max_iter=500)

Default behaviour: zipmap=True

The output type for the probabilities is a list of dictionaries.

sess = rt.InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
res = sess.run(None, {"X": X_test})
print(res[1][:2])
print("probabilities type:", type(res[1]))
print("type for the first observations:", type(res[1][0]))

[{10: 6.341787866404047e-06, 12: 0.030411386862397194, 14: 0.9695823192596436}, {10: 0.022245943546295166, 12: 0.9420960545539856, 14: 0.035658009350299835}]
probabilities type: <class 'list'>
type for the first observations: <class 'dict'>

Option zipmap=False

Probabilities are now a matrix.

initial_type = [("float_input", FloatTensorType([None, 4]))]
options = {id(clr): {"zipmap": False}}
onx2 = to_onnx(clr, X_train, options=options, target_opset=12)

sess2 = rt.InferenceSession(
    onx2.SerializeToString(), providers=["CPUExecutionProvider"]
)
res2 = sess2.run(None, {"X": X_test})
print(res2[1][:2])
print("probabilities type:", type(res2[1]))
print("type for the first observations:", type(res2[1][0]))

[[6.3417879e-06 3.0411387e-02 9.6958232e-01]
 [2.2245944e-02 9.4209605e-01 3.5658009e-02]]
probabilities type: <class 'numpy.ndarray'>
type for the first observations: <class 'numpy.ndarray'>

Option zipmap=’columns’

This options removes the final operator ZipMap and splits the probabilities into columns. The final model produces one output for the label, and one output per class.

options = {id(clr): {"zipmap": "columns"}}
onx3 = to_onnx(clr, X_train, options=options, target_opset=12)

sess3 = rt.InferenceSession(
    onx3.SerializeToString(), providers=["CPUExecutionProvider"]
)
res3 = sess3.run(None, {"X": X_test})
for i, out in enumerate(sess3.get_outputs()):
    print(
        "output: '{}' shape={} values={}...".format(
            out.name, res3[i].shape, res3[i][:2]
        )
    )

output: 'output_label' shape=(38,) values=[14 12]...
output: 'i10' shape=(38,) values=[6.3417879e-06 2.2245944e-02]...
output: 'i12' shape=(38,) values=[0.03041139 0.94209605]...
output: 'i14' shape=(38,) values=[0.9695823  0.03565801]...

Let’s compare prediction time

print("Average time with ZipMap:")
print(sum(repeat(lambda: sess.run(None, {"X": X_test}), number=100, repeat=10)) / 10)

print("Average time without ZipMap:")
print(sum(repeat(lambda: sess2.run(None, {"X": X_test}), number=100, repeat=10)) / 10)

print("Average time without ZipMap but with columns:")
print(sum(repeat(lambda: sess3.run(None, {"X": X_test}), number=100, repeat=10)) / 10)

# The prediction is much faster without ZipMap
# on this example.
# The optimisation is even faster when the classes
# are described with strings and not integers
# as the final result (list of dictionaries) may copy
# many times the same information with onnxruntime.

Average time with ZipMap:
0.003264389200558071
Average time without ZipMap:
0.0027052922996517736
Average time without ZipMap but with columns:
0.0026047332994494354

Option zimpap=False and output_class_labels=True

Option zipmap=False seems a better choice because it is much faster but labels are lost in the process. Option output_class_labels can be used to expose the labels as a third output.

initial_type = [("float_input", FloatTensorType([None, 4]))]
options = {id(clr): {"zipmap": False, "output_class_labels": True}}
onx4 = to_onnx(clr, X_train, options=options, target_opset=12)

sess4 = rt.InferenceSession(
    onx4.SerializeToString(), providers=["CPUExecutionProvider"]
)
res4 = sess4.run(None, {"X": X_test})
print(res4[1][:2])
print("probabilities type:", type(res4[1]))
print("class labels:", res4[2])

[[6.3417879e-06 3.0411387e-02 9.6958232e-01]
 [2.2245944e-02 9.4209605e-01 3.5658009e-02]]
probabilities type: <class 'numpy.ndarray'>
class labels: [10 12 14]

Processing time.

print("Average time without ZipMap but with output_class_labels:")
print(sum(repeat(lambda: sess4.run(None, {"X": X_test}), number=100, repeat=10)) / 10)

Average time without ZipMap but with output_class_labels:
0.003510921600536676

MultiOutputClassifier

This model is equivalent to several classifiers, one for every label to predict. Instead of returning a matrix of probabilities, it returns a sequence of matrices. Let’s first modify the labels to get a problem for a MultiOutputClassifier.

y = numpy.vstack([y, y + 100]).T
y[::5, 1] = 1000  # Let's a fourth class.
print(y[:5])

[[  10 1000]
 [  10  110]
 [  10  110]
 [  10  110]
 [  10  110]]

Let’s train a MultiOutputClassifier.

X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = MultiOutputClassifier(LogisticRegression(max_iter=500))
clr.fit(X_train, y_train)
print(clr)

onx5 = to_onnx(clr, X_train, target_opset=12)

sess5 = rt.InferenceSession(
    onx5.SerializeToString(), providers=["CPUExecutionProvider"]
)
res5 = sess5.run(None, {"X": X_test[:3]})
print(res5)

MultiOutputClassifier(estimator=LogisticRegression(max_iter=500))
/home/xadupre/github/sklearn-onnx/skl2onnx/_parse.py:564: UserWarning: Option zipmap is ignored for model <class 'sklearn.multioutput.MultiOutputClassifier'>. Set option zipmap to False to remove this message.
  warnings.warn(
[array([[ 10, 110],
       [ 12, 112],
       [ 12, 114]], dtype=int64), [array([[9.6185070e-01, 3.8149051e-02, 2.1035541e-07],
       [4.4240355e-02, 9.4647479e-01, 9.2849005e-03],
       [8.0552045e-03, 6.3393027e-01, 3.5801452e-01]], dtype=float32), array([[7.9632753e-01, 8.7586209e-02, 1.8680112e-04, 1.1589945e-01],
       [2.6829399e-02, 7.2040278e-01, 5.2288104e-02, 2.0047970e-01],
       [7.8932270e-03, 3.6448562e-01, 5.4698211e-01, 8.0639035e-02]],
      dtype=float32)]]

Option zipmap is ignored. Labels are missing but they can be added back as a third output.

onx6 = to_onnx(
    clr,
    X_train,
    target_opset=12,
    options={"zipmap": False, "output_class_labels": True},
)

sess6 = rt.InferenceSession(
    onx6.SerializeToString(), providers=["CPUExecutionProvider"]
)
res6 = sess6.run(None, {"X": X_test[:3]})
print("predicted labels", res6[0])
print("predicted probabilies", res6[1])
print("class labels", res6[2])

predicted labels [[ 10 110]
 [ 12 112]
 [ 12 114]]
predicted probabilies [array([[9.6185070e-01, 3.8149051e-02, 2.1035541e-07],
       [4.4240355e-02, 9.4647479e-01, 9.2849005e-03],
       [8.0552045e-03, 6.3393027e-01, 3.5801452e-01]], dtype=float32), array([[7.9632753e-01, 8.7586209e-02, 1.8680112e-04, 1.1589945e-01],
       [2.6829399e-02, 7.2040278e-01, 5.2288104e-02, 2.0047970e-01],
       [7.8932270e-03, 3.6448562e-01, 5.4698211e-01, 8.0639035e-02]],
      dtype=float32)]
class labels [array([10, 12, 14], dtype=int64), array([ 110,  112,  114, 1000], dtype=int64)]

Versions used for this example

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