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.
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
The directory, not only the .pb file, is required for full restoration.
2. Load as Keras model
If model was exported with Keras-compatible metadata, this is simplest path.
3. Load as generic SavedModel signature
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.
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.
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.pbfile instead of model directory. - Expecting TF1 frozen-graph loaders to work with TF2 Keras APIs.
- Losing custom layers/objects and forgetting
custom_objectsduring 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
- How to load a model from an HDF5 file in Keras?
- How to load a tflite model in script?
- How to load a trained model''s weights, which were saved with tf.keras.models.save_model?
- How to load a trained TF1 protobuf model into TF2?
- How to load an image and show the image using keras?
- how to load and use a saved model on tensorflow?
- How to load batches of CSV files using tf.data and map
- How to load Image Masks Labels for Image Segmentation in Keras
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.