TensorFlow
Keras
Machine Learning
Prediction Models
Deep Learning

Make predictions using a tensorflow graph from a keras model

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

In TensorFlow 2, you normally predict with model(x, training=False) or model.predict. If you specifically want to run inference from the TensorFlow graph representation instead of the full Keras object, the modern path is to export the model as a SavedModel and load its serving signature.

What "using the graph" means now

Older TensorFlow code often talked about grabbing tensors from a graph and running them in a session. In current TensorFlow, that is no longer the usual workflow. A Keras model can be exported as a graph-backed SavedModel, and tf.saved_model.load returns callable signatures for inference.

That means you do not usually extract raw nodes manually. You export the model, load the serving function, and call it with tensors.

Export and load a graph-backed serving function

The following example builds a small Keras model, exports it, reloads it through tf.saved_model.load, and uses the serving signature for prediction.

python
1import numpy as np
2import tensorflow as tf
3
4model = tf.keras.Sequential(
5    [
6        tf.keras.layers.Input(shape=(3,)),
7        tf.keras.layers.Dense(4, activation="relu"),
8        tf.keras.layers.Dense(1),
9    ]
10)
11
12model(np.zeros((1, 3), dtype=np.float32))
13
14tf.saved_model.save(model, "saved_model_dir")
15
16loaded = tf.saved_model.load("saved_model_dir")
17serve = loaded.signatures["serving_default"]
18
19x = tf.constant([[1.0, 2.0, 3.0]], dtype=tf.float32)
20outputs = serve(x)
21
22print(outputs)

The key point is that serve is graph-backed. You are no longer using the Keras predict method. You are calling the exported TensorFlow serving function directly.

When to use this instead of load_model

If you just want to reload the model in Python and keep the Keras API, tf.keras.models.load_model is usually simpler. Use the graph-backed approach when:

  • you are serving the model outside normal Keras training code
  • you want the exported serving signature
  • you are integrating with systems that consume SavedModel artifacts

Once you load with tf.saved_model.load, you are working with callable signatures and TensorFlow functions rather than a full compiled Keras model object.

Input signatures and outputs

The serving function has a defined input and output structure. That is why passing the right dtype and shape matters. In simple single-input models, calling serve(x) is often enough. In models with named inputs, the signature may expect keyword arguments keyed by the exported input names.

Inspecting loaded.signatures is a good habit because it tells you what callable endpoints are available. The common default is serving_default.

It is also useful to inspect the returned dictionary keys after one inference call. SavedModel signatures often return named outputs rather than a plain tensor, especially when the original model has multiple outputs or explicitly named layers. Treating the signature like a small serving API rather than like raw Keras internals helps avoid confusion.

That is also why this approach fits deployment boundaries well. The exported signature defines the public inference contract explicitly, so another process can call it without needing your original model-building code, optimizer state, or training configuration.

Common Pitfalls

  • Reaching for old session-and-graph patterns in TensorFlow 2 when SavedModel signatures are the intended inference API.
  • Forgetting to build the Keras model before saving it.
  • Expecting tf.saved_model.load to return a normal Keras model with predict, fit, and compile.
  • Passing NumPy arrays or tensors with the wrong shape or dtype to the serving signature.
  • Using graph export when load_model would have been simpler for pure Python reuse.

Summary

  • In TensorFlow 2, graph-based inference from a Keras model is usually done through SavedModel.
  • Export the model, load it with tf.saved_model.load, and call serving_default.
  • This gives you a graph-backed callable, not the full Keras training interface.
  • Use load_model when you want a normal Keras object back in Python.
  • Treat the serving signature as the stable prediction entry point for exported inference.

Course illustration
Course illustration

All Rights Reserved.