Introduction To Custom Training Loops

Keras's model.fit() handles training for you in one line, but it hides everything that happens underneath. Once you move into research-style projects, multi-loss models, or custom architectures like GANs, you need to take the wheel yourself. Custom training loops give you that control. This chapter unpacks exactly what fit() does internally and shows you how to rebuild it, piece by piece.

Introduction to Custom Training

A custom training loop is simply the manual version of model.fit(). Instead of TensorFlow looping over batches and updating weights for you, you write that loop yourself. This is not busywork for its own sake. Custom loops are essential when:

• Your model has multiple outputs or multiple loss functions that need to be combined in a non-standard way.

• You are training adversarial setups like GANs, where two networks update against each other in alternating steps.

• You need fine-grained control over what happens every single batch, such as logging gradients, applying custom regularisation, or skipping batches conditionally.

• You are implementing a research paper that uses a training procedure fit() simply cannot express.

The trade-off is clear: fit() is fast to write and well-optimised, but a custom loop is fully transparent. Every step, from prediction to weight update, is visible and editable. Once you understand the loop, debugging training issues becomes far easier because you can inspect any intermediate value.

Using tf.GradientTape

tf.GradientTape is the engine behind every custom training loop. It records every operation performed on a tensor while inside its context manager, building a computation graph on the fly. When you call tape.gradient(), TensorFlow walks that recorded graph backwards and computes the derivative of your loss with respect to every trainable variable.

import tensorflow as tf

x = tf.Variable(3.0)

with tf.GradientTape() as tape:

y = x ** 2 + 2 * x + 1   # y = x^2 + 2x + 1

dy_dx = tape.gradient(y, x)

print(dy_dx.numpy())   # 8.0  ->  dy/dx = 2x + 2 = 2(3)+2

Only tf.Variable objects (and tensors you explicitly tell the tape to watch) are tracked by default. Constants are ignored unless you call tape.watch() on them. Inside a training step, the pattern always follows the same three moves: forward pass inside the tape, compute the loss, then pull the gradients out.

model = tf.keras.Sequential([tf.keras.layers.Dense(10, activation='softmax')])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

with tf.GradientTape() as tape:

predictions = model(x_batch, training=True)

loss = loss_fn(y_batch, predictions)

gradients = tape.gradient(loss, model.trainable_variables)

Did You Know?

By default, a GradientTape discards its recorded operations the moment you call .gradient() once. If you need to compute multiple gradients from the same forward pass (common in GAN training), create the tape with persistent=True, and remember to delete it manually afterwards to free memory.

Writing a Training Loop from Scratch

With GradientTape as the gradient engine, a full training loop is just four repeated steps: forward pass, loss calculation, backward pass, and weight update, wrapped inside two for loops — one for epochs, one for batches.

import tensorflow as tf

model = tf.keras.Sequential([

tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),

tf.keras.layers.Dense(10, activation='softmax'),

])

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()

train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))

train_dataset = train_dataset.shuffle(1024).batch(64)

epochs = 5

for epoch in range(epochs):

print(f'Epoch {epoch + 1}/{epochs}')

for step, (x_batch, y_batch) in enumerate(train_dataset):

with tf.GradientTape() as tape:

logits = model(x_batch, training=True)

loss_value = loss_fn(y_batch, logits)

grads = tape.gradient(loss_value, model.trainable_variables)

optimizer.apply_gradients(zip(grads, model.trainable_variables))

train_acc_metric.update_state(y_batch, logits)

if step % 200 == 0:

print(f'  step {step}: loss = {loss_value:.4f}')

print(f'  training accuracy: {train_acc_metric.result():.4f}')

train_acc_metric.reset_state()

Notice that nothing here is magic. tf.data.Dataset handles shuffling and batching, the tape records the forward pass, apply_gradients() performs the weight update, and a tf.keras.metrics object tracks accuracy across the epoch. Every one of these pieces is something you can swap out or extend.

Gradient Clipping

Sometimes gradients explode, growing so large that a single weight update overshoots wildly and destabilises training. This is especially common in recurrent networks and very deep models. Gradient clipping caps the magnitude of gradients before they are applied, keeping training stable without changing the overall learning dynamics.

There are two common strategies:

•        Clip by value: each individual gradient component is clamped to a fixed range, e.g. between -1 and 1.

•        Clip by global norm: the entire gradient vector is rescaled so that its overall norm does not exceed a threshold, preserving the direction of the update while shrinking its size.

with tf.GradientTape() as tape:

logits = model(x_batch, training=True)

loss_value = loss_fn(y_batch, logits)

grads = tape.gradient(loss_value, model.trainable_variables)

# Clip by global norm — the most commonly used approach

grads, _ = tf.clip_by_global_norm(grads, clip_norm=1.0)

optimizer.apply_gradients(zip(grads, model.trainable_variables))

Clipping by global norm is generally preferred over clipping by value because it preserves the relative scale between different gradients, only shrinking the overall step when it is genuinely too large.

Gradient Accumulation

Gradient accumulation lets you simulate a large batch size when your GPU memory cannot hold one. Instead of updating weights after every batch, you sum (accumulate) gradients across several smaller batches, then apply the combined update once. The effective batch size becomes batch_size × accumulation_steps.

accum_steps = 4

accumulated_grads = [tf.zeros_like(v) for v in model.trainable_variables]

for step, (x_batch, y_batch) in enumerate(train_dataset):

with tf.GradientTape() as tape:

logits = model(x_batch, training=True)

loss_value = loss_fn(y_batch, logits) / accum_steps

grads = tape.gradient(loss_value, model.trainable_variables)

accumulated_grads = [ag + g for ag, g in zip(accumulated_grads, grads)]

if (step + 1) % accum_steps == 0:

optimizer.apply_gradients(zip(accumulated_grads, model.trainable_variables))

accumulated_grads = [tf.zeros_like(v) for v in model.trainable_variables]

Dividing the loss by accum_steps before computing gradients keeps the scale of the accumulated update equivalent to what a single large batch would have produced. This technique is widely used when fine-tuning large transformer models on limited GPU memory.

Quick Tip

Gradient accumulation trades wall-clock speed for memory: you still process the same number of samples, but across more, smaller forward and backward passes. Use it only when you are genuinely memory-constrained, not as a default.