C.W.K.
Stream
Lesson 01 of 07 · published

When to Subclass

~8 min · subclass

Level 0Keras Apprentice
0 XP0/97 lessons0/20 achievements
0/120 XP to next level120 XP to go0% complete

When even the graph isn't enough — when you need conditional control flow inside call(), dynamic shapes, or a custom training step — you subclass keras.Model or keras.layers.Layer. This is the *I'll just write Python* exit hatch. The cost: tooling that auto-introspects (model.summary(), plot_model) loses some X-ray vision on subclassed bits.

The escape hatch, and what it costs

The Sequential and Functional APIs describe a model as a static graph — a fixed wiring of layers that Keras builds before any data flows. That graph is what makes model.summary(), shape inference, and plot_model() work: Keras can read the whole topology without ever running it. Subclassing trades that graph for raw Python. You write call() like any method, and the forward pass becomes whatever your code does — including things a static graph can't express.

Reach for it when:

  • The forward pass branches on the input itself — tree-structured nets, mixture-of-experts routing, anything where the data decides the path
  • You need real Python control flow (if/else, loops over a variable count) inside the forward pass
  • The architecture is a moving research target and you want to edit logic, not rewire a graph
  • The computation simply isn't a static DAG — recurrent state machines, custom autoregressive decoding

The cost, stated plainly

A subclassed model can't be reconstructed from a config dict alone, can't be inspected as a graph until it has been called once on real data, and won't hand you intermediate activations as cleanly as a Functional model. So the question isn't "which API is better" — it's "how much of my model genuinely needs Python?"

The recommended pattern resolves the tension: subclass layers for the dynamic pieces, then compose those layers with the Functional API for the static skeleton. You pay the introspection cost only on the few blocks that earn it, and the rest of the model stays X-ray-visible. The next lessons build exactly this — first the raw subclass mechanics, then the seam where the two APIs meet.

External links

Exercise

List three architectures from recent papers (2023-2024) that you can express with Functional, and three you can't (need Subclass). Justify each.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.