Store arrays in one onnx graph#

Once a model is converted it can be useful to store an array as a constant in the graph an retrieve it through an output. This allows the user to store training parameters or other informations like a vocabulary. Last sections shows how to remove an output or to promote an intermediate result to an output.

Train and convert a model#

We download one model from the :epkg:`ONNX Zoo` but the model could be trained and produced by another converter library.

import pprint
import numpy
from onnx import load
from onnxruntime import InferenceSession
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from skl2onnx import to_onnx
from skl2onnx.helpers.onnx_helper import (
    add_output_initializer, select_model_inputs_outputs)


data = load_iris()
X, y = data.data.astype(numpy.float32), data.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
model = LogisticRegression(penalty='elasticnet', C=2.,
                           solver='saga', l1_ratio=0.5)
model.fit(X_train, y_train)

onx = to_onnx(model, X_train[:1], target_opset=12,
              options={'zipmap': False})

Out:

D:\Program Files\Python\Python39\lib\site-packages\sklearn\linear_model\_sag.py:352: ConvergenceWarning: The max_iter was reached which means the coef_ did not converge
  warnings.warn(

Add training parameter#

new_onx = add_output_initializer(
    onx,
    ['C', 'l1_ratio'],
    [numpy.array([model.C]), numpy.array([model.l1_ratio])])

Inference#

sess = InferenceSession(new_onx.SerializeToString())
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {'X': X_test[:2]})
print("outputs")
pprint.pprint(res)

Out:

output names: ['label', 'probabilities', 'C', 'l1_ratio']
outputs
[array([0, 2], dtype=int64),
 array([[9.9800849e-01, 1.9915334e-03, 3.5919868e-11],
       [1.1456419e-04, 1.7301913e-01, 8.2686633e-01]], dtype=float32),
 array([2.]),
 array([0.5])]

The major draw back of this solution is increase the prediction time as onnxruntime copies the constants for every prediction. It is possible either to store those constant in a separate ONNX graph or to removes them.

Select outputs#

Next function removes unneeded outputs from a model, not only the constants. Next model only keeps the probabilities.

simple_onx = select_model_inputs_outputs(new_onx, ['probabilities'])

sess = InferenceSession(simple_onx.SerializeToString())
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {'X': X_test[:2]})
print("outputs")
pprint.pprint(res)

# Function *select_model_inputs_outputs* add also promote an intermediate
# result to an output.
#

Out:

output names: ['probabilities']
outputs
[array([[9.9800849e-01, 1.9915334e-03, 3.5919868e-11],
       [1.1456419e-04, 1.7301913e-01, 8.2686633e-01]], dtype=float32)]

This example only uses ONNX graph in memory and never saves or loads a model. This can be done by using the following snippets of code.

Save a model#

with open("simplified_model.onnx", "wb") as f:
    f.write(simple_onx.SerializeToString())

Load a model#

model = load("simplified_model.onnx", "wb")

sess = InferenceSession(model.SerializeToString())
print("output names:", [o.name for o in sess.get_outputs()])
res = sess.run(None, {'X': X_test[:2]})
print("outputs")
pprint.pprint(res)

Out:

output names: ['probabilities']
outputs
[array([[9.9800849e-01, 1.9915334e-03, 3.5919868e-11],
       [1.1456419e-04, 1.7301913e-01, 8.2686633e-01]], dtype=float32)]

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

Gallery generated by Sphinx-Gallery