keras
model loading
TensorFlow
.pb format
machine learning

How to load a keras model saved as .pb

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

A Keras model saved in TensorFlow SavedModel format includes a saved_model.pb file plus variable directories. You do not load the .pb file alone directly in most Keras workflows; instead you load the SavedModel directory.

Many loading errors come from mixing old TensorFlow 1 graph-loading code with modern TensorFlow 2 Keras APIs. This guide shows correct loading methods and how to inspect signatures for inference.

Core Sections

1. Save in SavedModel format

python
model.save("./exported_model")  # creates saved_model.pb + variables/

The directory, not only the .pb file, is required for full restoration.

2. Load as Keras model

python
1import tensorflow as tf
2
3loaded = tf.keras.models.load_model("./exported_model")
4probs = loaded.predict(x_test)

If model was exported with Keras-compatible metadata, this is simplest path.

3. Load as generic SavedModel signature

python
1import tensorflow as tf
2
3sm = tf.saved_model.load("./exported_model")
4print(list(sm.signatures.keys()))
5infer = sm.signatures["serving_default"]
6out = infer(tf.constant(x_batch))

Use this when you need low-level serving signatures or non-Keras exports.

4. Compatibility and migration notes

If model originated from legacy TF1 graph freezing, conversion may be required before Keras loading works. Pin TensorFlow version in inference environment and add smoke tests with known input/output pairs.

5. Build a repeatable validation checklist

After implementing Keras SavedModel loading, create a small validation pack that runs the same way on developer machines, CI, and staging. The checklist should include a baseline case, an edge case, and a failure-path case with expected outcomes written in plain language. This avoids the common situation where a workflow appears correct in one environment but fails under a slightly different runtime, dependency version, or input distribution.

A useful checklist should also capture environment assumptions explicitly: runtime version, dependency versions, configuration flags, and external services required by the scenario. Teams often skip this because it feels obvious during initial implementation, but those hidden assumptions are exactly what cause regressions during upgrades and handoffs.

text
1validation checklist
2- baseline scenario with expected output shape and values
3- edge scenario with constrained or unusual input
4- failure scenario with expected fallback or error behavior
5- runtime/dependency/config assumptions for reproducibility

Treat this checklist as a versioned artifact. If code behavior changes, update expected results in the same pull request rather than relying on informal tribal memory. Coupling implementation and validation updates keeps Keras SavedModel loading reliable as the codebase evolves.

6. Operational hardening and maintenance

Long-term reliability for Keras SavedModel loading depends on observability and clear ownership. Add structured logs and metrics around the most failure-prone operations so incident responders can quickly identify whether failures come from input quality, configuration mismatch, external dependency drift, or code regressions. Without those signals, teams spend most of incident time reconstructing context instead of fixing root causes.

Also define who owns periodic compatibility checks. Libraries, runtimes, cloud APIs, and tooling change over time, and silent drift is common. Schedule lightweight smoke checks that run even when no feature work is active, and record results so there is an audit trail for when behavior started to diverge.

bash
# example maintenance check command pattern
make smoke-test

Finally, document rollback criteria ahead of time. If a deployment changes Keras SavedModel loading behavior unexpectedly, the team should know when to roll back immediately versus when to hot-fix forward. This turns operational response from improvisation into a controlled process and prevents repeated incidents.

Common Pitfalls

  • Passing path to saved_model.pb file instead of model directory.
  • Expecting TF1 frozen-graph loaders to work with TF2 Keras APIs.
  • Losing custom layers/objects and forgetting custom_objects during load.
  • Ignoring signature names when using tf.saved_model.load.
  • Loading model successfully but using mismatched preprocessing pipeline.

Summary

To load a Keras model saved as .pb within SavedModel, load the containing directory via Keras or SavedModel APIs. Choose Keras loader for standard inference, and signature-based loading for serving-level control. Version pinning and preprocessing parity are essential for correct predictions after restore.


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.