7 A deep dive on Keras
This chapter deepens your Keras expertise through the lens of “progressive disclosure of complexity,” showing how the same API supports a smooth spectrum of workflows from simple to highly flexible. It contrasts three ways to build models—Sequential for straightforward stacks, the Functional API for graph-like, multi-input/multi-output architectures with access to layer connectivity, and Model subclassing for full control at the cost of fewer built-ins and more responsibility. Crucially, these approaches interoperate seamlessly, letting you mix and match Functional models with subclassed layers or vice versa, and encouraging you to choose the right tool for each job.
On the training side, the chapter revisits the standard compile/fit/evaluate/predict loop and demonstrates how to enrich it with custom metrics, callbacks, and monitoring. You learn to implement metrics by subclassing Metric, track state across batches, and plug them into fit(). Callbacks such as EarlyStopping and ModelCheckpoint help automate robust training, while TensorBoard provides live visual feedback on metrics, model topology, and learned representations. Together, these tools turn a basic training run into a controlled, observable, and reproducible experiment loop.
When the built-in loop isn’t enough, the chapter shows how to write custom training logic—from single-step updates to full loops—across TensorFlow, PyTorch, and JAX. It covers training vs inference modes, the distinction between trainable and non-trainable weights, and backend-specific gradient handling: GradientTape and apply in TensorFlow, backward/zeroing grads in PyTorch, and JAX’s stateless paradigm using stateless_call and optimizer.stateless_apply. Finally, it bridges flexibility and convenience by overriding Model.train_step() to run bespoke algorithms while retaining fit() features (callbacks, distribution, performance), and explains metric handling in this setup for each backend.
Progressive disclosure of complexity for model building
Plot generated by plot_model() on our ticket classifier model
Model plot with shape information added
Plot of our new model
The output of our custom history plotting callback
The loop of progress HTML: class=“small-image”
TensorBoard can be used for easy monitoring of training and evaluation metrics
Chapter summary
- Keras offers a spectrum of different workflows, based on the principle of progressive disclosure of complexity. They all smoothly interoperate.
- You can build models via the
Sequentialclass, via the Functional API, or by subclassing theModelclass. Most of the time, you’ll be using the Functional API. - The simplest way to train and evaluate a model is via the default
fit()andevaluate()methods. - Keras callbacks provide a simple way to monitor models during your call to
fit()and automatically take action based on the state of the model. - You can also fully control what
fit()does by overriding thetrain_step()method, leveraging APIs from your backend of choice – JAX, TensorFlow, or PyTorch. - Beyond
fit(), you can also write your own training loops entirely from scratch, in a backend-native way. This is useful for researchers implementing brand new training algorithms.
FAQ
What does “progressive disclosure of complexity” mean in Keras?
Keras is designed so you can start simple and add complexity only when you need it. You can begin with easy, high-level workflows (like Sequential + fit) and progressively move to more flexible ones (Functional API, subclassing, custom training loops) without switching tools. All parts interoperate because they share the same Layer and Model abstractions.What are the three ways to build models in Keras, and how do they differ?
- Sequential: simplest “stack of layers,” one input and one output, linear topology.- Functional API: defines arbitrary directed acyclic graphs of layers (multi-input/output, branching, merging). Best trade-off between ease and flexibility.
- Model subclassing: write your own Model class and call method; maximal control, but fewer built-ins (no graph topology to inspect) and easier to make mistakes.
How do I pick between Sequential, Functional API, and subclassing?
- Use Sequential for quick prototypes with simple stacks.- Prefer the Functional API for most real projects: it handles complex topologies, supports plotting/feature extraction, and remains concise.
- Use subclassing only when your forward pass can’t be expressed as a DAG (loops, recursion, dynamic control flow). Be aware you’ll lose topology inspection and some tooling.
How do I ensure a Sequential model is “built” and inspect shapes as I add layers?
- Layers create weights on first call or when you call model.build(input_shape=...).- To get summaries while incrementally adding layers, insert an Input first: model.add(keras.Input(shape=(...))). Then model.summary() shows evolving output shapes. You can also name models/layers for clearer debugging.
How do I build and train a multi-input, multi-output model with the Functional API?
- Define one Input per input stream; combine via layers like Concatenate; branch into multiple output heads.- Train by passing lists in the same order as Model(inputs=..., outputs=...) or use dicts keyed by input/output names.
- In compile, provide loss/metrics as lists (ordered) or dicts (by name) to match each output.
How can I visualize model topology and reuse intermediate features?
- Use keras.utils.plot_model(model, ...) to render the graph (optionally show_shapes=True, show_layer_names=True).- Access layer connectivity to do feature extraction: via model.layers or by grabbing layer.output to create new models that reuse intermediate representations.
What are callbacks, and when should I use EarlyStopping and ModelCheckpoint?
- Callbacks hook into fit() to observe and control training.- EarlyStopping: stop when monitored metric stops improving (with patience).
- ModelCheckpoint: save weights/models during training (often save_best_only=True).
- Other common callbacks: LearningRateScheduler, ReduceLROnPlateau, CSVLogger, TensorBoard.
How do I write a custom metric in Keras?
Create a subclass of keras.metrics.Metric, define state variables in __init__ via add_weight, implement update_state(y_true, y_pred), result() to compute the metric value, and reset_state() to clear state between epochs. Then pass an instance to compile(metrics=[...]).How do I write a custom training loop, and what differs across TensorFlow, PyTorch, and JAX?
- Core steps: forward pass (model(inputs, training=True)), compute loss, compute gradients, apply optimizer updates to trainable_weights.- TensorFlow: use tf.GradientTape().
- PyTorch: loss.backward(), read v.value.grad, apply updates in a no_grad block, then zero grads.
- JAX: use stateless functions; model.stateless_call(...) and optimizer.stateless_apply(...); jax.value_and_grad with has_aux=True to carry non-trainable updates.
How can I customize model.fit() by overriding train_step(), and how do I handle metrics (incl. JAX)?
- Subclass keras.Model and override train_step(self, data). Inside, run forward, compute loss (self.compute_loss uses compile() loss + regularization), compute/apply grads, update metrics, and return a dict of metric names to values.- TensorFlow/PyTorch: update metrics with metric.update_state(...).
- JAX: use stateless metric APIs (metric.stateless_update_state/ stateless_result) and carry metric variables in the training state alongside model and optimizer variables.
Deep Learning with Python, Third Edition ebook for free
