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, 2], dtype=int64), array([[0.03960225, 0.79471505, 0.16568264],
       [0.00120798, 0.34261844, 0.65617365]], 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, 2], dtype=int64), array([[0.03960225, 0.79471505, 0.16568264],
       [0.00120798, 0.34261844, 0.65617365]], 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, 2], dtype=int64), array([[0.03960225, 0.79471505, 0.16568264],
       [0.00120798, 0.34261844, 0.65617365]], 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, 2], dtype=int64), array([[0.03960225, 0.79471505, 0.16568264],
       [0.00120798, 0.34261844, 0.65617365]], dtype=float32)]