Introduction

Graphing a tf.keras model is useful for architecture review, debugging shape mismatches, and documenting experiments. In TensorFlow 2, teams usually combine static diagrams, textual summaries, and TensorBoard traces to get full visibility. A robust setup includes dependency checks and fallbacks so visualization works in local machines and CI pipelines.

Quick Architectural View with model.summary

Start with summary because it has no external rendering dependencies.

1import tensorflow as tf
2from tensorflow.keras import layers
3
4model = tf.keras.Sequential([
5    layers.Input(shape=(32,)),
6    layers.Dense(64, activation="relu"),
7    layers.Dense(32, activation="relu"),
8    layers.Dense(1, activation="sigmoid")
9])
10
11model.summary()

This prints layer order, output shapes, and parameter counts. It is usually the fastest way to catch obvious architecture mistakes.

Static Diagram with plot_model

For PR reviews and docs, a rendered diagram is easier to scan than console text.

1import tensorflow as tf
2from tensorflow.keras.utils import plot_model
3
4plot_model(
5    model,
6    to_file="model_architecture.png",
7    show_shapes=True,
8    show_dtype=True,
9    show_layer_names=True,
10    expand_nested=True
11)

This generates an image you can attach to experiment artifacts.

Dependency note:

• 'plot_model usually requires pydot and Graphviz binaries.'
• In container builds, install both explicitly.

Example install commands:

pip install pydot
a pt-get update && apt-get install -y graphviz

If image generation fails, keep model.summary output as fallback artifact.

Functional API Example for Branching Graphs

Branching models are easier to verify visually than by reading code.

1import tensorflow as tf
2from tensorflow.keras import layers
3
4inputs = tf.keras.Input(shape=(128,), name="input")
5left = layers.Dense(64, activation="relu", name="left_dense")(inputs)
6right = layers.Dense(64, activation="relu", name="right_dense")(inputs)
7merged = layers.Concatenate(name="merge")([left, right])
8outputs = layers.Dense(10, activation="softmax", name="classifier")(merged)
9
10branch_model = tf.keras.Model(inputs=inputs, outputs=outputs, name="branch_model")
11branch_model.summary()

Run plot_model on this model to verify merge paths and output dimensions before training.

TensorBoard Graph Trace

Static diagrams show architecture, but TensorBoard can show traced execution graphs.

1import tensorflow as tf
2from datetime import datetime
3
4logdir = f"logs/graph/{datetime.now().strftime('%Y%m%d-%H%M%S')}"
5writer = tf.summary.create_file_writer(logdir)
6
7@tf.function
8def forward(x):
9    return model(x)
10
11sample = tf.random.uniform((1, 32))
12with writer.as_default():
13    tf.summary.trace_on(graph=True, profiler=False)
14    _ = forward(sample)
15    tf.summary.trace_export(name="model_trace", step=0)
16
17print("logdir:", logdir)

Then run:

tensorboard --logdir logs/graph

Open the Graph tab to inspect traced operations.

Subclassed Model Gotcha

Subclassed models often need a build step before graphing because shape metadata may be unknown.

1import tensorflow as tf
2from tensorflow.keras import layers
3
4class MyModel(tf.keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.d1 = layers.Dense(32, activation="relu")
8        self.d2 = layers.Dense(1)
9
10    def call(self, inputs):
11        x = self.d1(inputs)
12        return self.d2(x)
13
14m = MyModel()
15_ = m(tf.random.uniform((1, 16)))  # build by calling once
16m.summary()

Without this call, summary or plot steps can fail or show incomplete metadata.

CI-Friendly Workflow

A practical pipeline step can generate both text and image artifacts:

1. Build model with sample input.
2. Save summary text file.
3. Attempt diagram render.
4. Mark render as warning, not hard failure, when Graphviz is missing.

Example summary export:

with open("model_summary.txt", "w", encoding="utf-8") as f:
    model.summary(print_fn=lambda line: f.write(line + "\n"))

This ensures architecture evidence exists even when image tooling is unavailable.

Common Pitfalls

• Expecting plot_model to work without Graphviz. Fix by installing Graphviz and pydot in runtime.
• Reviewing only summary text for complex branching models. Fix by adding rendered diagrams for structural checks.
• Forgetting to build subclassed models before visualization. Fix by calling model once with sample input.
• Exporting TensorBoard traces before executing traced function. Fix by running at least one forward pass after trace_on.
• Treating visualization as optional during refactors. Fix by making graph artifacts part of model-review workflow.

Summary

• Use model.summary first for fast, dependency-light architecture inspection.
• Use plot_model to generate readable diagrams with shapes and layer names.
• Use TensorBoard trace export for execution-graph inspection.
• Build subclassed models before plotting or summary generation.
• Automate visualization artifacts in CI to catch architecture regressions early.