TensorFlow / Keras#

This page answers common “how do I…” questions for converting TensorFlow/Keras models to ONNX with yobx.tensorflow.to_onnx().

How to convert a TensorFlow model#

Build a Keras model, run one forward pass to initialize weights, then call yobx.tensorflow.to_onnx() with a representative input array:

<<<

import numpy as np
import tensorflow as tf
from yobx.helpers.onnx_helper import pretty_onnx
from yobx.tensorflow import to_onnx

model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(8, activation="relu", input_shape=(4,)),
        tf.keras.layers.Dense(3),
    ]
)

rng = np.random.default_rng(0)
X = rng.standard_normal((5, 4)).astype(np.float32)
_ = model(X)  # initialize model weights

onx = to_onnx(model, (X,))
print(pretty_onnx(onx))

>>>

    /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/keras/src/layers/core/dense.py:107: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
      super().__init__(activity_regularizer=activity_regularizer, **kwargs)
    opset: domain='' version=21
    input: name='X:0' type=dtype('float32') shape=['dim', 4]
    init: name='sequential_2/dense_4/kernel:0[sequential_2/dense_4/kernel_371]' type=float32 shape=(4, 8)-- _convert_concrete_function.0
    init: name='sequential_2/dense_5/kernel:0[sequential_2/dense_5/kernel_393]' type=float32 shape=(8, 3)-- _convert_concrete_function.0
    MatMul(X:0, sequential_2/dense_4/kernel:0[sequential_2/dense_4/kernel_371]) -> sequential_2_1/dense_4_1/MatMul:0
      Relu(sequential_2_1/dense_4_1/MatMul:0) -> sequential_2_1/dense_4_1/Relu:0
        MatMul(sequential_2_1/dense_4_1/Relu:0, sequential_2/dense_5/kernel:0[sequential_2/dense_5/kernel_393]) -> Identity:0
    output: name='Identity:0' type='NOTENSOR' shape=None

How to validate ONNX outputs with onnxruntime#

Run the exported model with onnxruntime and compare outputs with the original Keras model:

<<<

import numpy as np
import onnxruntime
import tensorflow as tf
from yobx.tensorflow import to_onnx

model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(16, activation="relu", input_shape=(6,)),
        tf.keras.layers.Dense(2),
    ]
)

rng = np.random.default_rng(0)
X_train = rng.standard_normal((20, 6)).astype(np.float32)
_ = model(X_train)  # initialize model weights
onx = to_onnx(model, (X_train[:1],))

X_test = rng.standard_normal((7, 6)).astype(np.float32)
expected = model(X_test).numpy()

sess = onnxruntime.InferenceSession(
    onx.SerializeToString(), providers=["CPUExecutionProvider"]
)
(got,) = sess.run(None, {"X:0": X_test})

assert np.allclose(expected, got, atol=1e-5)
print("Outputs match ✓")

>>>

    /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/keras/src/layers/core/dense.py:107: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
      super().__init__(activity_regularizer=activity_regularizer, **kwargs)
    Outputs match ✓

How to control dynamic shapes#

By default, the first input axis is dynamic. Pass dynamic_shapes to name it explicitly:

<<<

import numpy as np
import tensorflow as tf
from yobx.tensorflow import to_onnx

model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(4, activation="relu", input_shape=(3,)),
        tf.keras.layers.Dense(1),
    ]
)

rng = np.random.default_rng(0)
X = rng.standard_normal((2, 3)).astype(np.float32)
_ = model(X)  # initialize model weights

onx = to_onnx(model, (X,), dynamic_shapes=({0: "batch"},))
dim0 = onx.graph.input[0].type.tensor_type.shape.dim[0]
print(f"dynamic dim name: {dim0.dim_param!r}")

>>>

    /opt/hostedtoolcache/Python/3.12.13/x64/lib/python3.12/site-packages/keras/src/layers/core/dense.py:107: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.
      super().__init__(activity_regularizer=activity_regularizer, **kwargs)
    dynamic dim name: 'dim'