TensorFlow
SavedModel
Machine Learning
Model Prediction
AI Frameworks

TensorFlow How to predict from a SavedModel?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

SavedModel is TensorFlow's standard export format for serving and reuse. If you can load the model directory, inspect its signatures, and feed tensors with the expected shape and dtype, prediction becomes straightforward.

Load the Model and Inspect Its Signatures

The most important detail is that a SavedModel is not just weights on disk. It also contains callable functions, often exposed under a serving_default signature. When you load the directory with tf.saved_model.load, TensorFlow returns an object with those exported functions attached.

python
1import tensorflow as tf
2
3model = tf.saved_model.load("exported-model")
4
5print(list(model.signatures.keys()))
6predict_fn = model.signatures["serving_default"]
7print(predict_fn.structured_input_signature)
8print(predict_fn.structured_outputs)

structured_input_signature tells you the input names, shapes, and dtypes. That matters because prediction errors usually come from passing NumPy arrays or tensors that do not match the exported signature.

If you are exporting a model yourself, this small example creates a simple network and saves it in a way that is easy to reload later:

python
1import tensorflow as tf
2import numpy as np
3
4keras_model = tf.keras.Sequential(
5    [
6        tf.keras.layers.Input(shape=(2,)),
7        tf.keras.layers.Dense(4, activation="relu"),
8        tf.keras.layers.Dense(1),
9    ]
10)
11
12keras_model.compile(optimizer="adam", loss="mse")
13
14x = np.array([[1.0, 2.0], [2.0, 1.0], [3.0, 4.0]], dtype=np.float32)
15y = np.array([[3.0], [2.0], [7.0]], dtype=np.float32)
16
17keras_model.fit(x, y, epochs=3, verbose=0)
18keras_model.export("exported-model")

After export, the directory contains metadata, variables, and any assets the model needs at inference time.

Run Prediction with the Exported Signature

Once the model is loaded, call the signature like a regular TensorFlow function. The key is to pass keyword arguments that match the exported input name. In many Keras exports, the input name is something generic such as keras_tensor, so inspecting the signature first is worth the effort.

python
1import tensorflow as tf
2
3loaded = tf.saved_model.load("exported-model")
4predict_fn = loaded.signatures["serving_default"]
5
6inputs = tf.constant([[5.0, 1.0], [1.5, 2.5]], dtype=tf.float32)
7result = predict_fn(keras_tensor=inputs)
8
9print(result)

The return value is usually a dictionary mapping output names to tensors. You can convert the tensor to NumPy if you are working in eager mode:

python
predictions = result["output_0"].numpy()
print(predictions)

If you do not know the argument name ahead of time, inspect it dynamically:

python
1_, keyword_args = predict_fn.structured_input_signature
2input_name = next(iter(keyword_args))
3
4batch = tf.constant([[10.0, 2.0]], dtype=tf.float32)
5output = predict_fn(**{input_name: batch})
6print(output)

This approach is especially useful when the model was exported by another team or generated in a notebook months ago.

Use the Keras Loading Path When Appropriate

Some TensorFlow models are better loaded with Keras APIs instead of the low-level SavedModel interface. If the export was created with model.save() in a format that Keras understands, tf.keras.models.load_model gives you back an object with the familiar predict method.

python
1import tensorflow as tf
2import numpy as np
3
4model = tf.keras.models.load_model("my-keras-model.keras")
5batch = np.array([[4.0, 2.0], [8.0, 1.0]], dtype=np.float32)
6
7predictions = model.predict(batch)
8print(predictions)

The rule of thumb is simple:

  • Use tf.saved_model.load when you are consuming a generic TensorFlow export.
  • Use tf.keras.models.load_model when you are loading a Keras model and want the higher-level API back.

Common Pitfalls

The most common failure is sending the wrong input name. SavedModel functions are strict about keyword names, so guessing often produces an error even when the tensor shape is correct.

Another common issue is dtype mismatch. A model trained on float32 inputs may reject int64 arrays or silently force conversions that hurt performance. Create tensors with an explicit dtype during inference.

Shape mismatches also show up frequently. If the signature expects a batch shaped like (None, 2), a one-dimensional tensor such as [5.0, 1.0] is not enough. Wrap the sample in an outer batch dimension.

Developers also confuse output names. A result dictionary may use names such as output_0 or a layer-derived label rather than predictions. Print structured_outputs once and use the actual key.

Finally, do not assume every export has a serving_default signature. Some models define custom signatures, so list the available keys before calling anything.

Summary

  • 'tf.saved_model.load is the standard way to load a TensorFlow SavedModel directory.'
  • Inspect signatures first so you know the exact input names, shapes, and dtypes.
  • Call the prediction function with keyword arguments that match the exported signature.
  • Expect a dictionary of output tensors and inspect the real output keys.
  • Use the Keras loading path only when the model was saved in a Keras-compatible format.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.