Practice Now
Keras
TensorFlow
Subclassing API
Model Plotting
Deep Learning

How do I plot a Keras/Tensorflow subclassing API model?

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
Deep LearningTensorFlow & KerasNetworking & APIs

Introduction

The Keras Subclassing API (tf.keras.Model subclass) provides maximum flexibility for custom architectures, but tf.keras.utils.plot_model() cannot automatically plot subclassed models because the computation graph is not built until the model is called with data. Unlike Sequential or Functional API models, subclassed models define their forward pass in Python code (call() method), which Keras cannot inspect statically. You must build the model with a concrete input shape first, and even then the plot shows limited detail compared to Functional models.

The Problem

python
1import tensorflow as tf
2
3class MyModel(tf.keras.Model):
4    def __init__(self):
5        super().__init__()
6        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
7        self.dropout = tf.keras.layers.Dropout(0.3)
8        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')
9
10    def call(self, inputs, training=False):
11        x = self.dense1(inputs)
12        x = self.dropout(x, training=training)
13        return self.dense2(x)
14
15model = MyModel()
16
17# This fails — model has no graph yet
18tf.keras.utils.plot_model(model, show_shapes=True)
19# ValueError: This model has not yet been built.

Solution 1: Build the Model First

Call the model with sample data or use model.build() to create the graph.

python
1model = MyModel()
2
3# Option A: Build with input shape
4model.build(input_shape=(None, 784))
5
6# Option B: Call with sample data
7import numpy as np
8sample = np.random.randn(1, 784).astype(np.float32)
9model(sample)
10
11# Now plot works (but shows limited layer info)
12tf.keras.utils.plot_model(
13    model,
14    to_file='model.png',
15    show_shapes=True,
16    show_layer_names=True,
17    expand_nested=True
18)

The resulting plot shows the model as a single block with input/output shapes but does not show individual layer connections like a Functional API plot does.

Solution 2: Create a Functional Equivalent for Plotting

Build a Functional API model that mirrors your subclassed model, then plot it.

python
1class MyModel(tf.keras.Model):
2    def __init__(self):
3        super().__init__()
4        self.dense1 = tf.keras.layers.Dense(128, activation='relu')
5        self.dropout = tf.keras.layers.Dropout(0.3)
6        self.dense2 = tf.keras.layers.Dense(10, activation='softmax')
7
8    def call(self, inputs, training=False):
9        x = self.dense1(inputs)
10        x = self.dropout(x, training=training)
11        return self.dense2(x)
12
13    def build_graph(self, input_shape):
14        """Create a Functional model for visualization."""
15        inputs = tf.keras.Input(shape=input_shape)
16        return tf.keras.Model(inputs=inputs, outputs=self.call(inputs))
17
18model = MyModel()
19plot_model = model.build_graph((784,))
20
21tf.keras.utils.plot_model(
22    plot_model,
23    to_file='model_detailed.png',
24    show_shapes=True,
25    show_layer_names=True,
26    show_layer_activations=True,
27    dpi=150
28)

This produces a detailed graph showing each layer, its shape, and connections.

Solution 3: Use model.summary() for Text Output

python
1model = MyModel()
2model.build(input_shape=(None, 784))
3
4model.summary()
5# Model: "my_model"
6# _________________________________________________________________
7#  Layer (type)                Output Shape              Param #
8# =================================================================
9#  dense (Dense)               (None, 128)               100480
10#  dropout (Dropout)           (None, 128)               0
11#  dense_1 (Dense)             (None, 10)                1290
12# =================================================================
13# Total params: 101,770
14# Trainable params: 101,770
15# Non-trainable params: 0

Solution 4: Visualize with TensorBoard

python
1import tensorflow as tf
2import datetime
3
4model = MyModel()
5sample = tf.random.normal((1, 784))
6
7log_dir = "logs/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
8writer = tf.summary.create_file_writer(log_dir)
9
10# Trace the model to create a graph
11tf.summary.trace_on(graph=True)
12model(sample)
13with writer.as_default():
14    tf.summary.trace_export(name="model_trace", step=0)
bash
tensorboard --logdir=logs
# Navigate to the "Graphs" tab in the browser

Prerequisite: Install Graphviz

plot_model requires Graphviz and pydot.

bash
1# macOS
2brew install graphviz
3
4# Ubuntu/Debian
5sudo apt install graphviz
6
7# Python package
8pip install pydot graphviz

Common Pitfalls

  • Plotting before building: plot_model() requires the model to be built (weights allocated). Call model.build(input_shape=(...)) or pass sample data through the model first. Without this, you get ValueError: This model has not yet been built.
  • Expecting detailed layer graphs from subclassed models: Even after building, plot_model on a subclassed model shows a simplified view (single block). To get detailed layer-by-layer graphs, create a Functional API equivalent using tf.keras.Input and tf.keras.Model.
  • Missing Graphviz installation: plot_model depends on the system-level Graphviz binary and the pydot Python package. Missing either produces ImportError: Failed to import pydot or FileNotFoundError: "dot" not found in path. Install both the system package and the pip package.
  • Training-mode layers in build_graph: Layers like Dropout and BatchNormalization behave differently during training vs inference. When creating a Functional equivalent, the training parameter in call() defaults to False, so the plot reflects inference topology. This is usually correct for visualization purposes.
  • Nested subclassed models: If a subclassed model contains other subclassed models, expand_nested=True may not fully expand them. Each nested model must also implement build_graph() for complete visualization.

Summary

  • Subclassed models require model.build() or a forward pass before plot_model() works
  • For detailed layer graphs, create a Functional API equivalent via a build_graph() method
  • Use model.summary() for quick text-based architecture overview
  • Use TensorBoard graph tracing for interactive visualization
  • Install both Graphviz (system) and pydot (pip) as prerequisites
Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the 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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.