Introduction

Loading a TensorFlow 1 .pb model in TensorFlow 2 depends on what kind of .pb file you actually have. A frozen graph can be imported for inference through compatibility APIs, while a SavedModel-based export has a different loading path. The important distinction is that a TF1 protobuf graph is usually an inference graph, not a modern Keras model you can fine-tune directly in TF2.

Start by Identifying the Artifact

People often say “protobuf model” when they really mean one of two things:

• a frozen graph stored in a .pb file
• a TF1 SavedModel directory that contains protobuf metadata

If you have a single .pb file, it is often a frozen graph. That matters because the easiest TF2 workflow is to import it for inference, not to pretend it is a native TF2 model object.

Load a Frozen Graph in TF2 for Inference

TensorFlow 2 still provides compatibility helpers for importing a TF1 graph definition.

1import tensorflow as tf
2
3
4def load_frozen_graph(pb_path):
5    with tf.io.gfile.GFile(pb_path, 'rb') as f:
6        graph_def = tf.compat.v1.GraphDef()
7        graph_def.ParseFromString(f.read())
8
9    def _imports():
10        tf.compat.v1.import_graph_def(graph_def, name='')
11
12    wrapped = tf.compat.v1.wrap_function(_imports, [])
13    return wrapped
14
15
16wrapped = load_frozen_graph('model.pb')
17print(wrapped.graph.as_graph_def().node[:5])

This gives you a callable graph wrapper and access to the graph structure.

Find the Input and Output Tensors

A frozen graph does not automatically tell you the friendly Keras-style signature you might want, so you usually need to inspect tensor names.

ops = wrapped.graph.get_operations()
for op in ops[:20]:
    print(op.name)

Once you know the tensor names, you can extract them:

input_tensor = wrapped.graph.get_tensor_by_name('input:0')
output_tensor = wrapped.graph.get_tensor_by_name('output:0')

Then build a callable function:

1infer = wrapped.prune(
2    feeds=input_tensor,
3    fetches=output_tensor
4)

Now you can run inference in TF2 code:

1import numpy as np
2
3sample = tf.constant(np.random.rand(1, 224, 224, 3), dtype=tf.float32)
4result = infer(sample)
5print(result)

SavedModel Is Different

If the original TF1 export is actually a SavedModel directory, use tf.saved_model.load instead of importing a raw graph definition.

1import tensorflow as tf
2
3model = tf.saved_model.load('saved_model_dir')
4print(list(model.signatures.keys()))

This is generally a cleaner path than working with a frozen graph, because SavedModel preserves signatures more explicitly.

What You Cannot Assume

Loading a TF1 .pb into TF2 does not automatically give you:

• a tf.keras.Model
• a trainable modern model object
• an easy fine-tuning workflow

A frozen graph is mostly an inference artifact. If the goal is retraining or clean TF2-native serving, you often need a conversion or reconstruction step, not just raw loading.

When tf.compat.v1 Is the Right Tool

It is normal to use tf.compat.v1 for this kind of migration. TF2 did not remove the need to interoperate with older artifacts, but it does make the boundary explicit.

A sensible migration path is often:

1. load the TF1 graph for inference
2. verify input and output behavior
3. wrap it for serving or comparison
4. rebuild or convert to a more native TF2 form if long-term maintenance matters

Rebuilding May Be Better Than Importing Forever

If the model is strategically important, treating the imported frozen graph as a permanent TF2 solution is often not ideal. It may be better to:

• recover the original checkpoint if available
• rebuild the architecture in TF2 or Keras
• load transferable weights where possible
• re-export as SavedModel

The import-based solution is excellent for inference continuity, but not always the best long-term engineering endpoint.

Common Pitfalls

A common mistake is assuming every .pb file is a SavedModel. A single protobuf graph file and a SavedModel export are different artifacts.

Another issue is expecting a loaded TF1 frozen graph to behave like a native Keras model with .fit() and high-level layers. Frozen graphs are usually inference-oriented.

Developers also often get stuck on missing tensor names. Inspect the graph operations first instead of guessing input and output names.

Finally, use tf.compat.v1 deliberately. Migration code is a compatibility bridge, not a sign that TF2 itself is broken.

Summary

• In TF2, a TF1 .pb file is usually loaded through compatibility APIs for inference.
• Use tf.compat.v1.import_graph_def and wrap_function for frozen graphs.
• Inspect tensor names and prune the graph to build a usable callable inference function.
• Use tf.saved_model.load instead if the artifact is actually a SavedModel directory.
• For long-term maintenance or retraining, rebuilding into a native TF2 representation is often the better path.