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=17)

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(f"inputs={input_names!r}, outputs={output_names!r}")
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X'], outputs=['label', 'probabilities']
[array([2, 0], dtype=int64), array([[2.9775689e-03, 4.1962689e-01, 5.7739550e-01],
       [7.9938823e-01, 2.0051695e-01, 9.4790288e-05]], 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=17)

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(f"inputs={input_names!r}, outputs={output_names!r}")
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X56'], outputs=['label', 'probabilities']
[array([2, 0], dtype=int64), array([[2.9775689e-03, 4.1962689e-01, 5.7739550e-01],
       [7.9938823e-01, 2.0051695e-01, 9.4790288e-05]], 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=17)

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(f"inputs={input_names!r}, outputs={output_names!r}")
print(sess.run(None, {input_names[0]: X_test[:2]}))
inputs=['X'], outputs=['L', 'P']
[array([2, 0], dtype=int64), array([[2.9775689e-03, 4.1962689e-01, 5.7739550e-01],
       [7.9938823e-01, 2.0051695e-01, 9.4790288e-05]], 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(f"changed {proposed_name!r} into {result!r}.")
    return result


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

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(f"inputs={input_names!r}, outputs={output_names!r}")
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([2, 0], dtype=int64), array([[2.9775689e-03, 4.1962689e-01, 5.7739550e-01],
       [7.9938823e-01, 2.0051695e-01, 9.4790288e-05]], dtype=float32)]

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

Gallery generated by Sphinx-Gallery