Modify the ONNX graph#

This example shows how to change the default ONNX graph such as renaming the inputs or outputs names.

Basic example#

import numpy
from onnxruntime import InferenceSession
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from skl2onnx.common.data_types import FloatTensorType, Int64TensorType
from skl2onnx import to_onnx

iris = load_iris()
X, y = iris.data, iris.target
X = X.astype(numpy.float32)
X_train, X_test, y_train, y_test = train_test_split(X, y)

clr = LogisticRegression(solver="liblinear")
clr.fit(X_train, y_train)


onx = to_onnx(clr, X, options={"zipmap": False}, target_opset=15)

sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
input_names = [i.name for i in sess.get_inputs()]
output_names = [o.name for o in sess.get_outputs()]
print("inputs=%r, outputs=%r" % (input_names, output_names))
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X'], outputs=['label', 'probabilities']
[array([1, 1], dtype=int64), array([[0.05684776, 0.8244833 , 0.11866891],
       [0.05001822, 0.73737246, 0.21260935]], dtype=float32)]

Changes the input names#

It is possible to change the input name by using the parameter initial_types. However, the user must specify the input types as well.

onx = to_onnx(
    clr,
    X,
    options={"zipmap": False},
    initial_types=[("X56", FloatTensorType([None, X.shape[1]]))],
    target_opset=15,
)

sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
input_names = [i.name for i in sess.get_inputs()]
output_names = [o.name for o in sess.get_outputs()]
print("inputs=%r, outputs=%r" % (input_names, output_names))
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X56'], outputs=['label', 'probabilities']
[array([1, 1], dtype=int64), array([[0.05684776, 0.8244833 , 0.11866891],
       [0.05001822, 0.73737246, 0.21260935]], dtype=float32)]

Changes the output names#

It is possible to change the input name by using the parameter final_types.

onx = to_onnx(
    clr,
    X,
    options={"zipmap": False},
    final_types=[("L", Int64TensorType([None])), ("P", FloatTensorType([None, 3]))],
    target_opset=15,
)

sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
input_names = [i.name for i in sess.get_inputs()]
output_names = [o.name for o in sess.get_outputs()]
print("inputs=%r, outputs=%r" % (input_names, output_names))
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X'], outputs=['L', 'P']
[array([1, 1], dtype=int64), array([[0.05684776, 0.8244833 , 0.11866891],
       [0.05001822, 0.73737246, 0.21260935]], dtype=float32)]

Renaming intermediate results#

It is possible to rename intermediate results by using a prefix or by using a function. The result will be post-processed in order to unique names. It does not impact the graph inputs or outputs.

def rename_results(proposed_name, existing_names):
    result = "_" + proposed_name.upper()
    while result in existing_names:
        result += "A"
    print("changed %r into %r." % (proposed_name, result))
    return result


onx = to_onnx(clr, X, options={"zipmap": False}, naming=rename_results, target_opset=15)

sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
input_names = [i.name for i in sess.get_inputs()]
output_names = [o.name for o in sess.get_outputs()]
print("inputs=%r, outputs=%r" % (input_names, output_names))
print(sess.run(None, {input_names[0]: X_test[:2]}))
changed 'SklearnLinearClassifier' into '_SKLEARNLINEARCLASSIFIER'.
changed 'label' into '_LABEL'.
changed 'probabilities' into '_PROBABILITIES'.
changed 'LinearClassifier' into '_LINEARCLASSIFIER'.
changed 'probability_tensor' into '_PROBABILITY_TENSOR'.
changed 'Normalizer' into '_NORMALIZER'.
inputs=['X'], outputs=['label', 'probabilities']
[array([1, 1], dtype=int64), array([[0.05684776, 0.8244833 , 0.11866891],
       [0.05001822, 0.73737246, 0.21260935]], dtype=float32)]

Total running time of the script: (0 minutes 0.065 seconds)

Gallery generated by Sphinx-Gallery