TensorFlow Custom Training Loops
Custom training loops give you complete control over every aspect of model training. While model.fit() handles the training loop automatically, some situations demand you write the loop yourself: training GANs with two competing models, implementing research papers with non-standard update rules, applying different learning rates to different parts of a model, or collecting custom statistics during training. Custom loops are also the gateway to understanding what TensorFlow does internally on every training step.
When to Use a Custom Training Loop
Use model.fit(): Use Custom Loop:
───────────────────────────────────────────────────────
Standard classification Training GANs
Standard regression Multiple interleaved models
Transfer learning Custom gradient manipulation
Most real-world projects Implementing research papers
Non-standard loss functions
Step-by-step debugging
The Four Core Components
A custom training loop requires four components working together:
- GradientTape — records operations so TensorFlow can compute gradients automatically
- Loss function — measures how wrong the model's prediction is
- Optimizer — uses gradients to update weights
- Metrics — track progress without affecting training
Understanding tf.GradientTape
GradientTape is the key tool for custom training. It "records" forward-pass operations on a "tape" (like a cassette tape recording audio). After the forward pass completes, you "play back" the tape in reverse to calculate how much each weight contributed to the loss. This playback process is automatic differentiation — the mathematical engine behind backpropagation.
import tensorflow as tf
x = tf.constant(3.0)
with tf.GradientTape() as tape:
tape.watch(x) # Track this variable
y = x ** 2 + 2 * x # y = x² + 2x
dy_dx = tape.gradient(y, x)
print(dy_dx) # dy/dx at x=3: 2(3) + 2 = 8.0
For model weights (which are tf.Variable objects), GradientTape tracks them automatically without needing tape.watch().
A Simple Custom Training Step
import tensorflow as tf
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
tf.keras.layers.Dense(1, activation='sigmoid')
])
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
loss_fn = tf.keras.losses.BinaryCrossentropy()
@tf.function # Compiles this function to a TensorFlow graph for speed
def train_step(x_batch, y_batch):
with tf.GradientTape() as tape:
# Forward pass
predictions = model(x_batch, training=True)
loss = loss_fn(y_batch, predictions)
# Compute gradients
gradients = tape.gradient(loss, model.trainable_variables)
# Apply gradients to update weights
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss
The @tf.function decorator compiles the function into an optimized TensorFlow computation graph. The first call takes slightly longer as TensorFlow traces the function. Every subsequent call runs much faster than pure Python because TensorFlow executes the compiled graph directly.
The Complete Custom Training Loop
import tensorflow as tf
import numpy as np
# Data
x_train = np.random.random((1000, 20)).astype('float32')
y_train = np.random.randint(0, 2, size=(1000, 1)).astype('float32')
x_val = np.random.random((200, 20)).astype('float32')
y_val = np.random.randint(0, 2, size=(200, 1)).astype('float32')
# Model
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid')
])
optimizer = tf.keras.optimizers.Adam(0.001)
loss_fn = tf.keras.losses.BinaryCrossentropy()
# Metrics
train_loss_metric = tf.keras.metrics.Mean(name='train_loss')
train_acc_metric = tf.keras.metrics.BinaryAccuracy(name='train_acc')
val_loss_metric = tf.keras.metrics.Mean(name='val_loss')
val_acc_metric = tf.keras.metrics.BinaryAccuracy(name='val_acc')
# Datasets
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(1000).batch(32)
val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(32)
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
preds = model(x, training=True)
loss = loss_fn(y, preds)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
train_loss_metric.update_state(loss)
train_acc_metric.update_state(y, preds)
@tf.function
def val_step(x, y):
preds = model(x, training=False)
loss = loss_fn(y, preds)
val_loss_metric.update_state(loss)
val_acc_metric.update_state(y, preds)
# Main training loop
for epoch in range(30):
# Reset metrics at the start of each epoch
train_loss_metric.reset_state()
train_acc_metric.reset_state()
val_loss_metric.reset_state()
val_acc_metric.reset_state()
# Training
for x_batch, y_batch in train_ds:
train_step(x_batch, y_batch)
# Validation
for x_batch, y_batch in val_ds:
val_step(x_batch, y_batch)
print(
f"Epoch {epoch+1:02d} | "
f"loss={train_loss_metric.result():.4f} | "
f"acc={train_acc_metric.result():.4f} | "
f"val_loss={val_loss_metric.result():.4f} | "
f"val_acc={val_acc_metric.result():.4f}"
)
Why training=True and training=False Matter
Some layers behave differently during training and inference:
- Dropout — randomly zeroes out neurons during training (for regularization). During inference, all neurons are active.
- BatchNormalization — uses batch statistics during training. During inference, it uses the running statistics computed across all training batches.
Passing training=True activates these training-specific behaviors. Passing training=False switches them to inference mode. Forgetting this distinction causes subtle bugs where your model performs worse during evaluation than it should.
Gradient Clipping: Preventing Exploding Gradients
In deep networks or RNNs, gradients can become very large during backpropagation, causing the model's weights to receive enormous updates that destabilize training. Gradient clipping caps the magnitude of gradients before applying them:
@tf.function
def train_step_with_clipping(x, y):
with tf.GradientTape() as tape:
preds = model(x, training=True)
loss = loss_fn(y, preds)
gradients = tape.gradient(loss, model.trainable_variables)
# Clip gradients so their total norm never exceeds 1.0
gradients, _ = tf.clip_by_global_norm(gradients, clip_norm=1.0)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss
Accumulating Gradients for Large Effective Batch Sizes
Sometimes the batch size you need for stable training does not fit in GPU memory. Gradient accumulation simulates a large batch by summing gradients over several small batches before applying the update:
accumulation_steps = 8 # Simulate batch_size × 8
accumulated_gradients = [tf.Variable(tf.zeros_like(v))
for v in model.trainable_variables]
for step, (x_batch, y_batch) in enumerate(train_ds):
with tf.GradientTape() as tape:
preds = model(x_batch, training=True)
loss = loss_fn(y_batch, preds) / accumulation_steps
grads = tape.gradient(loss, model.trainable_variables)
for acc_grad, grad in zip(accumulated_gradients, grads):
acc_grad.assign_add(grad)
if (step + 1) % accumulation_steps == 0:
optimizer.apply_gradients(
zip(accumulated_gradients, model.trainable_variables)
)
# Reset accumulated gradients
for acc_grad in accumulated_gradients:
acc_grad.assign(tf.zeros_like(acc_grad))
Saving Checkpoints Inside the Loop
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
checkpoint_manager = tf.train.CheckpointManager(
checkpoint, directory='./checkpoints', max_to_keep=3
)
for epoch in range(30):
# ... training steps ...
# Save checkpoint every 5 epochs
if (epoch + 1) % 5 == 0:
save_path = checkpoint_manager.save()
print(f"Checkpoint saved: {save_path}")
Custom Loop Diagram
For each epoch: │ ├── For each batch in training data: │ │ │ ├── [GradientTape context starts recording] │ │ │ ├── Forward Pass: x → model → predictions │ │ │ ├── Compute Loss: loss_fn(y_true, predictions) │ │ │ ├── [GradientTape stops recording] │ │ │ ├── Backprop: tape.gradient(loss, trainable_vars) │ │ │ ├── (Optional) Clip gradients │ │ │ └── optimizer.apply_gradients(grads, vars) │ ├── For each batch in validation data: │ └── model(x, training=False) → compute val metrics │ └── Print epoch summary → repeat
Custom training loops trade convenience for control. Once you understand this pattern, you can implement any training procedure described in a research paper. The next topic covers Callbacks — objects that plug into both model.fit() and custom loops to automate actions like saving the best model, adjusting the learning rate, or stopping training early when the model stops improving.