TensorFlow
freezing model
protobuf
TensorFlow 2
machine learning

Freezing graph to pb in Tensorflow2

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

In TensorFlow 2, "freezing a graph" no longer means working with the old TensorFlow 1 session workflow. The modern path is usually to export a SavedModel, but if you specifically need a frozen .pb graph for inference tooling or interoperability, you can still generate one from a concrete function.

What Freezing Means in TensorFlow 2

A frozen graph is a graph definition in which variable values are converted into constants. That makes the graph self-contained for inference because the weights are embedded directly into the graph representation.

In TensorFlow 2, the usual deployment artifact is SavedModel, not a frozen graph. That distinction matters:

  • 'SavedModel is the standard TensorFlow export format'
  • a frozen .pb graph is a more specialized artifact
  • many older tutorials assume TensorFlow 1 graph sessions and no longer apply directly

If your target tool accepts SavedModel, prefer that first.

The Modern Freezing Flow

The common TensorFlow 2 workflow is:

  1. build or load a Keras model
  2. wrap the model call in a tf.function
  3. get a concrete function with explicit input shape
  4. convert variables to constants
  5. write the resulting graph to a .pb file

Here is a small working example:

python
1import tensorflow as tf
2from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(4,)),
6    tf.keras.layers.Dense(8, activation="relu"),
7    tf.keras.layers.Dense(3)
8])
9
10model(tf.zeros((1, 4)))
11
12full_model = tf.function(lambda x: model(x))
13concrete = full_model.get_concrete_function(
14    tf.TensorSpec([None, 4], tf.float32)
15)
16
17frozen_func = convert_variables_to_constants_v2(concrete)
18frozen_graph = frozen_func.graph.as_graph_def()
19
20tf.io.write_graph(
21    graph_or_graph_def=frozen_graph,
22    logdir="./export",
23    name="model_frozen.pb",
24    as_text=False,
25)

This writes a binary .pb file that contains the frozen inference graph.

Why the Input Signature Matters

The input signature you pass to get_concrete_function determines the graph shape and input types captured in the exported artifact. If the signature is too narrow or wrong, the frozen graph may not match your inference workload.

For example, an image model might need:

python
tf.TensorSpec([None, 224, 224, 3], tf.float32)

Choosing the signature carefully is part of the export contract.

Inspecting Inputs and Outputs

After freezing, print the graph endpoints so you know what to feed at inference time:

python
1print("Inputs:")
2for tensor in frozen_func.inputs:
3    print(tensor)
4
5print("Outputs:")
6for tensor in frozen_func.outputs:
7    print(tensor)

This is important because the tensor names used by external inference tools come from the frozen graph, not from your original training code abstractions.

When You Should Not Freeze

If your goal is TensorFlow Serving, Keras reloading, or ordinary model deployment within TensorFlow tooling, SavedModel is usually better. It preserves more structure and fits the modern ecosystem.

Frozen graphs are most useful when:

  • a downstream runtime explicitly expects .pb
  • you need a self-contained graph for inference-only use
  • you are interoperating with older or specialized tooling

The fact that freezing is still possible does not mean it is the default deployment path.

Common Pitfalls

The biggest mistake is following TensorFlow 1 tutorials that use sessions and graph_util directly. Those patterns do not map cleanly to TensorFlow 2 eager execution.

Another common issue is forgetting to build the model before exporting. If the model has not been called, the graph may not be fully traced.

Developers also often freeze the graph without recording the input and output tensor names. That makes later inference integration much harder.

Finally, do not assume a frozen .pb is the best export format just because it sounds portable. In many TensorFlow 2 deployments, SavedModel is the better answer.

Summary

  • In TensorFlow 2, export SavedModel by default and freeze only when you specifically need a .pb graph.
  • Freezing means converting variables into constants inside a concrete function graph.
  • Use convert_variables_to_constants_v2 on a traced tf.function.
  • Define the input signature carefully because it shapes the exported graph.
  • Record graph inputs and outputs so the frozen artifact is actually usable for inference.

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.