Module 3 — Exporting from TensorFlow

Course 08 finished with a Keras model saved as a SavedModel. Course 09 finished with a PyTorch state dict. Both need to reach the same destination — an ONNX file consumable by ONNX Runtime — but through different tools. Where PyTorch ships the exporter inside the framework, TensorFlow relies on a separate community package, tf2onnx, maintained by the ONNX organization. This module walks through the two supported entry points (SavedModel and Keras H5), the signature choice, and the recurring convention flip that trips up every team that mixes the two frameworks: NHWC versus NCHW.

The path is straightforward once you know where the sharp edges are. Getting there without knowing them costs a full day of debugging a model that appears to run but produces meaningless predictions.

Two ways in: SavedModel and Keras

The recommended input to tf2onnx is a SavedModel directory, produced by tf.saved_model.save. It carries a well-defined signature — inputs and outputs with names, shapes and dtypes — and is the format TensorFlow itself uses for serving.

pip install tf2onnx onnx onnxruntime

import tensorflow as tf

model = tf.keras.models.load_model("plant_disease_mobilenet_v2.h5")
tf.saved_model.save(model, "saved_model_dir")

From the command line, one call produces the .onnx file:

python -m tf2onnx.convert \
    --saved-model saved_model_dir \
    --output plant_disease.onnx \
    --opset 17

The same conversion is available in Python, which is easier to script and to combine with a numerical parity check afterwards:

import tf2onnx

spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=17,
    output_path="plant_disease.onnx",
)

input_signature plays the same role as example in PyTorch: it fixes dtype, dimensionality, and — through None — which axes are dynamic. The first axis marked None becomes the batch dimension; nothing else needs declaring, which is a real convenience compared to torch.onnx.export.

Signatures and the "concrete function" trap

A Keras model has one obvious signature: the call method taking inputs and returning outputs. A SavedModel can carry several — the training signature, the inference signature, a preprocessing signature — each stored as a concrete function. If you save a model without being explicit, TensorFlow picks defaults that may include intermediate tensors, batch normalisation statistics, or other artefacts you did not intend to export.

@tf.function(input_signature=[tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input")])
def serve(images):
    return {"logits": model(images, training=False)}

tf.saved_model.save(model, "saved_model_dir", signatures={"serving_default": serve})

The wrapper accomplishes two things at once. It fixes the input shape and dtype, and it passes training=False explicitly, which is the TensorFlow equivalent of PyTorch's model.eval(). Skipping this call is the direct cause of the same bug — dropout and batch normalisation staying in training mode — and produces the same symptom: a served model whose outputs drift from the trained one.

NHWC versus NCHW: the layout that flips

TensorFlow stores 2D convolutions as NHWC: batch, height, width, channels. PyTorch, ONNX and most GPU inference stacks use NCHW: batch, channels, height, width. Same numbers, different order. On a training image of (224, 224, 3), the tensor shapes differ:

Framework	Convention	Shape of one image	Shape of a batch of 32
TensorFlow / Keras	NHWC	(224, 224, 3)	(32, 224, 224, 3)
PyTorch, ONNX	NCHW	(3, 224, 224)	(32, 3, 224, 224)

tf2onnx handles the conversion automatically for standard layers: it inserts a Transpose node at the input, changes the operator attributes, and inserts another Transpose at the output if needed. The exported ONNX file still expects NHWC input by default, matching the TensorFlow convention. The team consuming the model in a PyTorch-style pipeline receives shape mismatches on the first request.

The clean fix is to tell tf2onnx to emit an NCHW-facing model with the --inputs-as-nchw flag:

python -m tf2onnx.convert \
    --saved-model saved_model_dir \
    --inputs-as-nchw input \
    --output plant_disease.onnx \
    --opset 17

Now the exported model takes (batch, 3, 224, 224), matching PyTorch clients. Every serving stack that mixes exports from both frameworks should standardise on one convention — usually NCHW — and enforce it at export time, not by adding Transpose calls in the serving code.

A Transpose inside the graph is normal; a Transpose in the client is a bug

When tf2onnx inserts a Transpose at the graph boundary, ONNX Runtime folds it into the first convolution during optimization (module 5) and the cost is often zero. When a Python client calls numpy.transpose before every request, the CPU pays for that permutation on every invocation, which can dominate the inference time on small models.

Opset and unsupported operators

TensorFlow has many more operators than the standard ONNX opset supports directly. tf2onnx maintains a mapping table; when it cannot map an op, it either decomposes it into primitives or fails with Unsupported ops. Two categories cause almost all such failures:

• String and vocabulary ops used in text preprocessing (tf.strings.lookup, HashTable) — these belong outside the model, not inside its graph.
• Custom ops compiled with tf.function and a Python for-loop that touches Python data structures — these need rewriting into pure tensor operations.

The standard remedy is to trim the model at the training-to-inference boundary. Preprocessing that uses string tables, text vectorisation or Python loops must live in the serving code, not in the ONNX graph. What crosses into ONNX is the numeric core: preprocessing tensors in, logits out.

Verifying the export before you leave TensorFlow

The same onnx.checker and numerical parity check that module 4 details apply verbatim. Before disconnecting the TensorFlow environment — where the source of truth still lives — run:

import onnx, onnxruntime as ort
import numpy as np, tensorflow as tf

model = tf.keras.models.load_model("plant_disease_mobilenet_v2.h5")
example_hwc = np.random.rand(1, 224, 224, 3).astype(np.float32)
example_chw = np.transpose(example_hwc, (0, 3, 1, 2))     # match --inputs-as-nchw

tf_out = model(example_hwc, training=False).numpy()

onnx.checker.check_model(onnx.load("plant_disease.onnx"))
sess = ort.InferenceSession("plant_disease.onnx", providers=["CPUExecutionProvider"])
onnx_out = sess.run(None, {"input": example_chw})[0]

print(f"max abs diff: {np.abs(tf_out - onnx_out).max():.2e}")

A difference around 1e-5 is normal for float32. A larger gap almost always traces to training=False missing, a preprocessing step that lived inside the Keras model, or an NHWC/NCHW mismatch that the check itself just revealed.

In summary

• tf2onnx converts a SavedModel or a Keras model into ONNX; the SavedModel path is preferred because it makes the signature explicit and reviewable.
• A serving signature wrapped in @tf.function with training=False fixes both the input contract and the equivalent of PyTorch's model.eval() mode.
• TensorFlow's NHWC convention differs from ONNX's NCHW; use --inputs-as-nchw to standardise at export time, not in every client.
• Unsupported ops almost always come from preprocessing that should live outside the model — string tables, vocabulary lookups, Python-side loops.

Next module: verifying numerical equivalence — the mandatory check that turns "the export succeeded" into "the export is correct".