advanced · ~30 min
Custom Training Loops
Drop `model.fit` and write the gradient tape by hand — inspect every gradient tensor as it flows backward.
model.fit is convenient because it hides an enormous amount of bookkeeping: computing gradients,
applying them, tracking metrics, running callbacks. Most of the time that's exactly what you want.
But sometimes you need to see — or change — what's happening inside that loop: a custom loss that
combines several models' outputs, a GAN's alternating generator/discriminator updates, gradient
inspection for debugging a model that refuses to train, or just full control for its own sake.
That's what a custom training loop gives you.
The core tool is tf.GradientTape (in
TensorFlow.js, the same idea shows up as optimizer.computeGradients): run the forward pass inside
it, and it remembers enough to hand back the gradient of your loss with respect to every trainable
variable, on request. From there, you decide what happens to those gradients before they're
applied — including not applying some of them at all, or shrinking ones that have gotten
dangerously large.
🔍 Deep dive: Why gradients explode, and what clipping actually does
A gradient's size roughly tracks how steep the loss surface is at the current point. Most of the time that's a small, well-behaved number — but occasionally, especially early in training or with a learning rate that's too aggressive for the model's current state, a single batch can produce a huge gradient. Left alone, one huge gradient step can throw the weights somewhere far worse than where they started — a spike in the loss curve, sometimes bad enough that training never recovers. Clipping by global norm rescales the entire gradient vector (all parameters together, preserving its direction) down to a maximum length whenever it exceeds a threshold, capping how far any single step can move the weights without changing what direction it moves them in.
The playground below runs two models side by side on the same hard dataset with a deliberately aggressive learning rate — one with gradient clipping on, one without — and charts each step's loss and its raw (pre-clipping) gradient norm, so you can watch a spike in gradient norm show up as a spike in loss for the unclipped model, and see clipping keep the clipped one from following it there.
Production note
In real training loops, tf.clip_by_global_norm(grads, clip_norm) is the standard call — it
clips the whole gradient list together (preserving relative direction across parameters), unlike
tf.clip_by_norm or tf.clip_by_value, which act on each tensor independently and can distort
the update direction. A clip norm of 1.0-5.0 is a common starting point; RNNs and transformers
reach for it more often than CNNs, since they're more prone to this kind of instability.
Beginner tip
If you're used to model.fit, the gradient-norm chart is worth staring at even without touching
the clipping toggle — it's the same information model.fit computes internally on every step,
just normally invisible to you.
Playground
Training loss
- clipped
- unclipped
Gradient norm (pre-clipping)
- clipped
- unclipped