Introduction

When working with TensorFlow, you often need to inspect or export the computational graph underlying your model. Whether you are debugging layer connections, optimizing for deployment, or converting between formats, extracting the Graph or GraphDef is a fundamental skill. This article walks you through the different approaches for TensorFlow 1.x, TensorFlow 2.x, and Keras models, so you can choose the right method for your situation.

Understanding Graph vs GraphDef

Before diving into extraction methods, it helps to understand why these two objects exist. A tf.Graph is TensorFlow's in-memory representation of a computational graph. It holds operations, tensors, and their relationships. A GraphDef is the serialized Protocol Buffer version of that graph, which you can save to disk, transfer between systems, or load into a different runtime. Think of Graph as the live object and GraphDef as its portable snapshot.

Extracting the Graph in TensorFlow 1.x

In TF1, every operation you create lands in a default graph, and sessions hold a reference to it. This makes extraction straightforward.

1import tensorflow as tf
2
3# Get the default graph
4graph = tf.get_default_graph()
5
6# If you have a session, get its graph
7with tf.Session() as sess:
8    graph = sess.graph
9
10    # Convert the graph to a GraphDef (serializable protobuf)
11    graph_def = graph.as_graph_def()
12
13    # Inspect operations
14    for op in graph.get_operations():
15        print(op.name, op.type)

The as_graph_def() method returns a GraphDef protobuf that captures every operation and tensor in the graph at that moment. This is the object you serialize when saving .pb files.

Extracting the Graph in TensorFlow 2.x with tf.function

TensorFlow 2.x defaults to eager execution, which means there is no global graph sitting around to grab. To get a graph, you need to trace a function using tf.function, which converts your Python code into a reusable computational graph.

1import tensorflow as tf
2
3@tf.function
4def my_computation(x):
5    return x * 2 + 1
6
7# Call the function once to trigger tracing
8concrete_func = my_computation.get_concrete_function(
9    tf.TensorSpec(shape=[None, 10], dtype=tf.float32)
10)
11
12# Access the graph and graph_def
13graph = concrete_func.graph
14graph_def = graph.as_graph_def()
15
16print(f"Number of operations: {len(graph.get_operations())}")

The key insight is that get_concrete_function forces TensorFlow to trace the function with a specific input signature, producing a ConcreteFunction whose .graph attribute gives you the traced graph.

Extracting the Graph from a Keras Model

For Keras models, wrap the model's call method with tf.function to generate a traceable graph.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
5    tf.keras.layers.Dense(10, activation='softmax')
6])
7
8# Wrap the model call in a tf.function
9@tf.function(input_signature=[
10    tf.TensorSpec(shape=[None, 784], dtype=tf.float32)
11])
12def model_forward(x):
13    return model(x, training=False)
14
15concrete_func = model_forward.get_concrete_function()
16graph_def = concrete_func.graph.as_graph_def()

Setting training=False ensures you capture the inference graph, which excludes dropout and batch normalization training behavior.

Freezing and Saving a GraphDef as .pb

A frozen graph bundles the model's weights directly into the GraphDef, producing a single self-contained file. This is essential for deployment scenarios where you want one portable artifact.

1from tensorflow.python.framework.convert_to_constants import (
2    convert_variables_to_constants_v2,
3)
4
5frozen_func = convert_variables_to_constants_v2(concrete_func)
6frozen_graph_def = frozen_func.graph.as_graph_def()
7
8# Save to a .pb file
9tf.io.write_graph(
10    frozen_graph_def,
11    logdir="./saved_model",
12    name="frozen_model.pb",
13    as_text=False
14)
15
16# Load it back
17with tf.io.gfile.GFile("./saved_model/frozen_model.pb", "rb") as f:
18    loaded_graph_def = tf.compat.v1.GraphDef()
19    loaded_graph_def.ParseFromString(f.read())

Using tf.compat.v1 for Legacy Code

If you are maintaining a codebase that mixes TF1 and TF2 patterns, tf.compat.v1 provides the bridge. Disabling eager execution restores the TF1 session-based workflow.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5# Build a graph using TF1-style code
6x = tf.compat.v1.placeholder(tf.float32, shape=[None, 10], name="input")
7y = tf.compat.v1.layers.dense(x, 5, name="output")
8
9with tf.compat.v1.Session() as sess:
10    sess.run(tf.compat.v1.global_variables_initializer())
11    graph_def = sess.graph.as_graph_def()

This approach is useful when you need to load older saved models or interface with tools that expect a TF1-style GraphDef.

Common Pitfalls

• Calling as_graph_def() in eager mode without tf.function leads to an empty or nonexistent graph because there is no graph to extract in pure eager execution.
• Forgetting to call get_concrete_function with an input signature means TensorFlow cannot trace the function and will raise an error about unknown shapes.
• Omitting training=False in Keras model wrapping captures training-only operations like dropout, producing an inference graph that behaves differently than expected.
• Confusing frozen graphs with SavedModel format results in lost flexibility, since a frozen .pb file converts variables to constants and cannot be fine-tuned afterward.
• Mixing tf.compat.v1.disable_eager_execution() with eager code is a global setting that affects all subsequent TensorFlow operations, causing subtle bugs in other parts of your script.

Summary

• Use sess.graph.as_graph_def() in TensorFlow 1.x to extract graphs from sessions.
• In TensorFlow 2.x, wrap your logic in tf.function and call get_concrete_function() to produce a traceable graph.
• For Keras models, wrap model(x, training=False) in a tf.function with an explicit input_signature.
• Freeze graphs with convert_variables_to_constants_v2 to bundle weights into a single .pb file for deployment.
• Use tf.compat.v1 when working with legacy TF1 codebases, but avoid mixing it with eager execution in the same script.
• Always provide explicit input shapes when tracing, since TensorFlow needs concrete dimensions to build the graph.