Introduction
TensorBoard is TensorFlow's visualization toolkit for inspecting model architecture, training metrics, and computational graphs. While TensorBoard is often used with summary operations (tf.summary) for logging scalars and histograms during training, you can also save just the computational graph for visualization without any summary operations. This is useful for inspecting model architecture, debugging graph structure, and documenting model design.
TensorFlow 2.x: Save Graph with tf.summary.trace
In TF 2.x with eager execution, use tf.summary.trace_export to capture the graph of a @tf.function:
1import tensorflow as tf
2import datetime
3
4# Build a model
5model = tf.keras.Sequential([
6 tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
7 tf.keras.layers.Dropout(0.2),
8 tf.keras.layers.Dense(64, activation='relu'),
9 tf.keras.layers.Dense(10, activation='softmax')
10])
11
12# Set up a log directory
13log_dir = "logs/graph/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
14writer = tf.summary.create_file_writer(log_dir)
15
16# Trace the graph
17tf.summary.trace_on(graph=True, profiler=False)
18
19# Run the model once to build the graph
20sample_input = tf.random.normal([1, 784])
21model(sample_input)
22
23# Export the trace
24with writer.as_default():
25 tf.summary.trace_export(name="model_trace", step=0)
26
27print(f"Graph saved to {log_dir}")
28print(f"View with: tensorboard --logdir {log_dir}")
TensorFlow 2.x: Save Graph with tf.function
For custom functions (not Keras models):
1import tensorflow as tf
2
3@tf.function
4def my_computation(x):
5 hidden = tf.nn.relu(tf.linalg.matvec(tf.Variable(tf.random.normal([128, 784])), x))
6 return tf.nn.softmax(hidden)
7
8# Trace the function
9log_dir = "logs/custom_graph"
10writer = tf.summary.create_file_writer(log_dir)
11
12tf.summary.trace_on(graph=True)
13result = my_computation(tf.random.normal([784]))
14with writer.as_default():
15 tf.summary.trace_export(name="custom_computation", step=0)
TensorFlow 1.x: Save Graph with FileWriter
In TF 1.x, the graph exists explicitly and can be saved directly:
1import tensorflow as tf
2
3# Build the graph
4x = tf.placeholder(tf.float32, shape=[None, 784], name='input')
5W1 = tf.Variable(tf.random_normal([784, 128]), name='weights_1')
6b1 = tf.Variable(tf.zeros([128]), name='bias_1')
7hidden = tf.nn.relu(tf.matmul(x, W1) + b1, name='hidden')
8
9W2 = tf.Variable(tf.random_normal([128, 10]), name='weights_2')
10b2 = tf.Variable(tf.zeros([10]), name='bias_2')
11output = tf.nn.softmax(tf.matmul(hidden, W2) + b2, name='output')
12
13# Save graph to TensorBoard log directory — no summary ops needed
14with tf.Session() as sess:
15 writer = tf.summary.FileWriter('logs/tf1_graph', sess.graph)
16 writer.close()
17
18# View: tensorboard --logdir logs/tf1_graph
The key line is tf.summary.FileWriter('logdir', sess.graph) — this writes the graph structure without requiring any tf.summary.scalar, tf.summary.histogram, or merge_all operations.
Save Keras Model Graph
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4 tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
5 tf.keras.layers.MaxPooling2D(),
6 tf.keras.layers.Flatten(),
7 tf.keras.layers.Dense(128, activation='relu'),
8 tf.keras.layers.Dense(10, activation='softmax')
9])
10
11# Method 1: Use TensorBoard callback (also logs training metrics)
12tensorboard_callback = tf.keras.callbacks.TensorBoard(
13 log_dir='logs/keras_graph',
14 write_graph=True, # Save the graph
15 write_images=False, # Don't save weight images
16)
17
18# Just call the model once — no training needed for graph only
19model.predict(tf.random.normal([1, 28, 28, 1]))
20
21# Method 2: Manual graph export
22log_dir = 'logs/keras_manual'
23writer = tf.summary.create_file_writer(log_dir)
24tf.summary.trace_on(graph=True)
25model(tf.random.normal([1, 28, 28, 1]))
26with writer.as_default():
27 tf.summary.trace_export(name="keras_model", step=0)
Save as GraphDef (Protocol Buffer)
Save the graph as a .pb file for inspection without TensorBoard:
1# TF 2.x
2import tensorflow as tf
3
4model = tf.keras.Sequential([
5 tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
6 tf.keras.layers.Dense(1)
7])
8
9# Get the concrete function
10concrete_func = tf.function(model).get_concrete_function(
11 tf.TensorSpec([None, 10], tf.float32)
12)
13
14# Save as SavedModel (includes graph)
15tf.saved_model.save(model, 'saved_model/')
16
17# Or export just the graph def
18graph_def = concrete_func.graph.as_graph_def()
19tf.io.write_graph(graph_def, '.', 'model_graph.pb', as_text=False)
20tf.io.write_graph(graph_def, '.', 'model_graph.pbtxt', as_text=True)
Viewing in TensorBoard
1# Install TensorBoard
2pip install tensorboard
3
4# Launch TensorBoard
5tensorboard --logdir logs/
6
7# Open in browser
8# http://localhost:6006
9
10# Specify a different port
11tensorboard --logdir logs/ --port 6007
12
13# View a specific run
14tensorboard --logdir logs/graph/20260301-120000
In TensorBoard, navigate to the Graphs tab to see the computational graph. You can:
Expand/collapse operation groups
View tensor shapes on edges
Search for specific operations
Color nodes by structure, device, or compute time
Common Pitfalls
Empty graph in TensorBoard: In TF 2.x with eager execution, you must use tf.summary.trace_on(graph=True) before running the model and tf.summary.trace_export() after. Without tracing, no graph is captured.
Graph not updating: TensorBoard caches data. If you re-export the graph, use a new log directory or restart TensorBoard with --reload_interval=5 for frequent updates.
Large graphs are hard to read: Complex models produce graphs with thousands of nodes. Use tf.name_scope('block_name') to group related operations, making the TensorBoard visualization navigable.
TF 1.x FileWriter must be closed: If you forget writer.close(), the events file may be incomplete or empty. Always close the writer or use it as a context manager.
SavedModel vs graph only: tf.saved_model.save() saves the graph plus weights and signatures. For visualization-only purposes, the tf.summary.trace_export approach is lighter.
Summary
TF 2.x: Use tf.summary.trace_on(graph=True) → run model → tf.summary.trace_export() to save the graph
TF 1.x: Use tf.summary.FileWriter('logdir', sess.graph) — no summary ops needed
Keras: Set write_graph=True in TensorBoard callback, or manually trace with tf.summary
View with tensorboard --logdir logs/ and navigate to the Graphs tab
Use tf.name_scope() to organize complex graphs into readable groups
