Tensorflow
Python
Model Restoration
.pb File
Machine Learning

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.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5graph = tf.Graph()
6with graph.as_default():
7    graph_def = tf.compat.v1.GraphDef()
8    with tf.io.gfile.GFile("model.pb", "rb") as f:
9        graph_def.ParseFromString(f.read())
10        tf.import_graph_def(graph_def, name="")
11
12with tf.compat.v1.Session(graph=graph) as sess:
13    ops = [op.name for op in graph.get_operations()[:10]]
14    print(ops)

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.

python
with tf.compat.v1.Session(graph=graph) as sess:
    input_tensor = graph.get_tensor_by_name("input:0")
    output_tensor = graph.get_tensor_by_name("output:0")

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.

python
1import numpy as np
2
3sample = np.random.randn(1, 224, 224, 3).astype("float32")
4
5with tf.compat.v1.Session(graph=graph) as sess:
6    result = sess.run(output_tensor, feed_dict={input_tensor: sample})
7    print(result.shape)

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:

python
for op in graph.get_operations()[:20]:
    print(op.name)

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 .pb file 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 .pb file often represents a graph artifact, commonly for inference.
  • Use tf.compat.v1.GraphDef and import_graph_def to 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.

Course illustration
Course illustration

All Rights Reserved.