How to restore Tensorflow model from .pb file in python?
Master System Design with Codemia
Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.
Introduction
Restoring a TensorFlow model from a .pb file depends on what that file actually contains. In older TensorFlow workflows, a .pb file often stores a graph definition, not necessarily the full training checkpoint state in the same way modern SavedModel formats do. The right restore method is therefore about identifying the file type first and then loading it with the correct TensorFlow API.
Understand What a .pb File Usually Contains
A .pb file commonly stores a serialized protocol buffer. In TensorFlow one style workflows, it often contains:
- a frozen graph for inference
- or a graph definition without the separate variable checkpoint files
That means you cannot assume every .pb file can be resumed for training directly. Many are inference-oriented artifacts.
Load a Frozen Graph in TensorFlow One Style
For graph-based inference, import the graph definition into a Graph.
This restores the computation graph into memory so you can inspect tensors and run inference.
Find Input and Output Tensor Names
After import, you usually need the input and output tensor names.
Those names depend on the original export. Inspecting operations is often necessary before inference works.
Run Inference from the Imported Graph
Once tensor names are known, feed data and fetch outputs.
This is the typical use case for a frozen .pb model.
Know the Difference from Checkpoint Restore
If you need to continue training, a frozen .pb alone is usually not enough. Training state typically requires:
- graph structure
- variable values
- optimizer state if resuming exactly
Those are more naturally handled by checkpoints or SavedModel exports, not just a frozen graph file.
In other words, restoring for inference is not the same as restoring for training.
TensorFlow Two Compatibility
In modern TensorFlow, the preferred model export format is SavedModel. If you only have a .pb frozen graph, you are usually operating in compatibility mode.
That means tf.compat.v1 is often the correct tool, even inside a TensorFlow two installation.
Trying to load old .pb graphs with only high-level Keras APIs usually leads to confusion because the artifact formats are different.
Inspect the Graph Before Assuming Names
Hard-coding tensor names from memory is a common failure. Print the operations when debugging:
This is often the fastest way to identify actual placeholders, outputs, and preprocessing nodes.
If You Can Re-Export, Prefer SavedModel
If you still control the original training code, re-exporting to SavedModel is usually a better long-term move than clinging to a raw frozen graph.
SavedModel gives:
- clearer loading APIs
- better TensorFlow two support
- better tooling compatibility
The .pb graph path is mainly a maintenance and interoperability workflow now.
Common Pitfalls
- Assuming every
.pbfile contains everything needed to resume training. - Guessing tensor names instead of inspecting the imported graph.
- Mixing TensorFlow one graph APIs and TensorFlow two eager expectations.
- Trying to use Keras model-loading functions on a frozen graph artifact.
- Forgetting to disable eager execution when using graph-session style code.
Summary
- A
.pbfile often represents a graph artifact, commonly for inference. - Use
tf.compat.v1.GraphDefandimport_graph_defto load it. - Inspect operations to discover real input and output tensor names.
- Use session-based execution for inference on legacy frozen graphs.
- Prefer SavedModel if you control the export path and need long-term maintainability.

