TensorFlow
neural network visualization
plotting neural networks
machine learning
Python programming

how to plot the tensorflow neural network object

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

If you want to "plot" a TensorFlow neural network, the right tool depends on what you mean by plotting. Sometimes you want a layer diagram, sometimes a text summary, and sometimes an interactive graph for debugging training behavior.

Use plot_model for a Layer Diagram

For Keras models, the most direct option is tf.keras.utils.plot_model.

python
1from tensorflow import keras
2
3model = keras.Sequential([
4    keras.layers.Input(shape=(32,)),
5    keras.layers.Dense(64, activation="relu"),
6    keras.layers.Dense(10, activation="softmax"),
7])
8
9keras.utils.plot_model(
10    model,
11    to_file="model.png",
12    show_shapes=True,
13    show_dtype=False,
14    show_layer_names=True,
15)

This generates a static image showing the layer structure. It is the easiest option for documentation and quick inspection.

Make Sure the Model Is Built First

Some models need to be built before plotting, especially subclassed models or models without an explicit input shape.

python
1from tensorflow import keras
2import tensorflow as tf
3
4class MyModel(keras.Model):
5    def __init__(self):
6        super().__init__()
7        self.dense1 = keras.layers.Dense(16, activation="relu")
8        self.dense2 = keras.layers.Dense(1)
9
10    def call(self, inputs):
11        x = self.dense1(inputs)
12        return self.dense2(x)
13
14model = MyModel()
15model(tf.zeros((1, 8)))  # build by calling once
16
17keras.utils.plot_model(model, to_file="subclassed_model.png", show_shapes=True)

If you skip the build step, plotting or summary output may be incomplete.

Use model.summary() for a Fast Text View

Sometimes a visual PNG is unnecessary. A text summary is often enough.

python
model.summary()

This gives:

  • layer names
  • output shapes
  • parameter counts

It is the fastest way to confirm that the network matches your intended architecture.

Use TensorBoard for Graph and Training Views

If you want more than a static diagram, TensorBoard is the better tool.

python
1import tensorflow as tf
2from tensorflow import keras
3
4model = keras.Sequential([
5    keras.layers.Input(shape=(16,)),
6    keras.layers.Dense(32, activation="relu"),
7    keras.layers.Dense(1),
8])
9
10model.compile(optimizer="adam", loss="mse")
11
12x = tf.random.normal((64, 16))
13y = tf.random.normal((64, 1))
14
15tb = keras.callbacks.TensorBoard(log_dir="./logs")
16model.fit(x, y, epochs=1, callbacks=[tb], verbose=0)

Then launch:

bash
tensorboard --logdir ./logs

TensorBoard is better when you want graph views plus metrics, not just a simple architecture snapshot.

Plotting Requirements

plot_model usually relies on Graphviz and pydot. If plotting fails even though your model is valid, the environment may be missing those dependencies.

Typical setup:

bash
pip install pydot

System Graphviz may also be required depending on the environment. If static plotting is unavailable, model.summary() still works and is often enough for debugging.

Functional Models Usually Plot Most Cleanly

Functional Keras models often produce the clearest diagrams because the input and output tensors are explicit.

python
1from tensorflow import keras
2
3inputs = keras.Input(shape=(8,))
4x = keras.layers.Dense(16, activation="relu")(inputs)
5outputs = keras.layers.Dense(1)(x)
6model = keras.Model(inputs, outputs)
7
8keras.utils.plot_model(model, to_file="functional_model.png", show_shapes=True)

If you are teaching, documenting, or reviewing architecture, this style is often easier to visualize than heavily dynamic subclassed code.

Which Tool to Choose

A practical rule is:

  • use model.summary() for quick inspection
  • use plot_model for a static architecture diagram
  • use TensorBoard for interactive graph and training visualization

These tools complement each other rather than replace one another.

Common Pitfalls

The biggest mistake is trying to plot a model before it is built, especially with subclassed models.

Another issue is assuming every TensorFlow object can be plotted the same way. Keras models are straightforward, but lower-level TensorFlow graphs and custom execution paths may need TensorBoard or other tooling instead.

A third problem is missing Graphviz-related dependencies and then assuming the model itself is broken.

Summary

  • Use tf.keras.utils.plot_model for a static image of a Keras network.
  • Use model.summary() for the fastest architectural inspection.
  • Build subclassed models before trying to plot them.
  • Use TensorBoard when you want richer graph and training visualization.
  • If plotting fails, check Graphviz and pydot dependencies before debugging the model itself.

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.