TensorFlow Callbacks

Callbacks are objects that run automatically at specific points during training — at the start of each epoch, at the end of each batch, when a metric improves, or when training finishes. They automate actions you would otherwise have to code manually: saving the best model, stopping early when learning plateaus, adjusting the learning rate, or logging metrics to a dashboard. Using the right callbacks turns a fragile training script into a robust, self-managing process.

How Callbacks Work

Training loop timeline:

on_train_begin()
│
├── Epoch 1
│     on_epoch_begin()
│     ├── Batch 1: on_batch_begin() → train step → on_batch_end()
│     ├── Batch 2: on_batch_begin() → train step → on_batch_end()
│     └── ...
│     on_epoch_end()  ← Most callbacks act here (save, check metrics, etc.)
│
├── Epoch 2  ...
│
on_train_end()
# Pass callbacks to model.fit as a list
model.fit(
    x_train, y_train,
    epochs=100,
    validation_data=(x_val, y_val),
    callbacks=[callback1, callback2, callback3]
)

EarlyStopping

EarlyStopping halts training when a monitored metric stops improving. This prevents wasting time on epochs that only make the model overfit, and it automatically restores the weights from the best epoch.

import tensorflow as tf

early_stop = tf.keras.callbacks.EarlyStopping(
    monitor='val_accuracy',   # Watch this metric
    patience=5,               # Stop after 5 epochs with no improvement
    mode='max',               # 'max' for accuracy; 'min' for loss
    restore_best_weights=True,# Revert to best epoch's weights when stopped
    min_delta=0.001           # Minimum change to count as improvement
)
Training with EarlyStopping(patience=5, monitor='val_accuracy'):

Epoch 1:  val_accuracy=0.70  ← best so far
Epoch 2:  val_accuracy=0.74  ← best so far
Epoch 3:  val_accuracy=0.77  ← best so far
Epoch 4:  val_accuracy=0.76  patience=1
Epoch 5:  val_accuracy=0.75  patience=2
Epoch 6:  val_accuracy=0.77  patience=3  (same as best — no improvement)
Epoch 7:  val_accuracy=0.76  patience=4
Epoch 8:  val_accuracy=0.76  patience=5  ← STOP. Restore weights from Epoch 3.

ModelCheckpoint

ModelCheckpoint saves the model to disk at regular intervals or only when the monitored metric improves. It protects against crashes and ensures you always have the best model saved.

checkpoint = tf.keras.callbacks.ModelCheckpoint(
    filepath='models/model_epoch_{epoch:02d}_val{val_accuracy:.3f}.keras',
    monitor='val_accuracy',
    save_best_only=True,   # Only save when val_accuracy improves
    mode='max',
    verbose=1
)

ReduceLROnPlateau

ReduceLROnPlateau automatically lowers the learning rate when the validation loss stops improving. A smaller learning rate helps the optimizer make finer adjustments to escape a local plateau in the loss surface.

reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.5,       # Multiply current LR by 0.5 when triggered
    patience=3,       # Wait 3 epochs before reducing
    min_lr=1e-7,      # Never go below this learning rate
    verbose=1
)
Epoch 1:  val_loss=0.80  LR=0.001
Epoch 5:  val_loss=0.52  LR=0.001
Epoch 6:  val_loss=0.53  patience=1
Epoch 7:  val_loss=0.52  patience=2
Epoch 8:  val_loss=0.53  patience=3  → LR reduced to 0.0005
Epoch 9:  val_loss=0.49  LR=0.0005  (improvement after LR drop)

TensorBoard

TensorBoard is a web-based visualization tool for monitoring training metrics, model architecture, and gradient histograms. The TensorBoard callback writes log files that the TensorBoard server reads and displays.

tensorboard_cb = tf.keras.callbacks.TensorBoard(
    log_dir='./logs',          # Directory for log files
    histogram_freq=1,          # Log weight histograms every epoch
    write_graph=True,          # Visualize model graph
    update_freq='epoch'        # Log metrics every epoch
)

# After training, launch TensorBoard:
# tensorboard --logdir ./logs
# Open http://localhost:6006 in your browser

CSVLogger

CSVLogger appends training metrics to a CSV file after each epoch. This creates a permanent record you can analyze with spreadsheet tools or pandas.

csv_logger = tf.keras.callbacks.CSVLogger(
    filename='training_log.csv',
    separator=',',
    append=False   # Set True to continue logging from a resumed training run
)

# training_log.csv content after 3 epochs:
# epoch,accuracy,loss,val_accuracy,val_loss
# 0,0.623,0.912,0.601,0.945
# 1,0.741,0.672,0.722,0.731
# 2,0.803,0.512,0.785,0.601

LearningRateScheduler

LearningRateScheduler changes the learning rate according to a custom function at the start of each epoch. This gives you full control over the learning rate schedule.

def lr_schedule(epoch, current_lr):
    # Warm up for 5 epochs, then decay every 10 epochs
    if epoch < 5:
        return 0.001 * (epoch + 1) / 5   # Linear warmup
    elif epoch % 10 == 0:
        return current_lr * 0.5           # Halve every 10 epochs
    return current_lr

lr_scheduler = tf.keras.callbacks.LearningRateScheduler(
    schedule=lr_schedule,
    verbose=1
)

Custom Callback

Build your own callback by subclassing tf.keras.callbacks.Callback and overriding the methods you need.

class MetricsLogger(tf.keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs=None):
        print(f"\n─── Starting Epoch {epoch + 1} ───")

    def on_epoch_end(self, epoch, logs=None):
        acc    = logs.get('accuracy', 0)
        val_acc = logs.get('val_accuracy', 0)
        lr     = float(self.model.optimizer.learning_rate)
        print(f"Epoch {epoch+1:03d} | "
              f"acc={acc:.4f} | val_acc={val_acc:.4f} | lr={lr:.6f}")

    def on_train_end(self, logs=None):
        print("\nTraining complete.")
        print(f"Final val_accuracy: {logs.get('val_accuracy', 0):.4f}")

custom_cb = MetricsLogger()

Combining Multiple Callbacks

import tensorflow as tf

model.fit(
    train_ds,
    epochs=100,
    validation_data=val_ds,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(
            monitor='val_accuracy', patience=8,
            restore_best_weights=True
        ),
        tf.keras.callbacks.ModelCheckpoint(
            'best_model.keras', monitor='val_accuracy',
            save_best_only=True
        ),
        tf.keras.callbacks.ReduceLROnPlateau(
            monitor='val_loss', factor=0.5, patience=4
        ),
        tf.keras.callbacks.TensorBoard(log_dir='./logs'),
        tf.keras.callbacks.CSVLogger('training_log.csv')
    ]
)

Callback Priority Summary

Callback              Use It When
──────────────────────────────────────────────────────────────────
EarlyStopping         Always — prevents wasted epochs and overfitting
ModelCheckpoint       Always — protects against crashes
ReduceLROnPlateau     Training plateaus; loss oscillates
TensorBoard           Monitoring training visually; debugging
CSVLogger             Permanent training record for analysis
LearningRateScheduler Custom warmup, cosine decay, or step schedules
Custom Callback       Any action not covered by built-in callbacks
──────────────────────────────────────────────────────────────────

Callbacks transform training from a manual monitoring task into an automatic, self-adjusting process. Combining EarlyStopping, ModelCheckpoint, and ReduceLROnPlateau gives you a robust training pipeline that saves the best model, avoids overfitting, and adapts the learning rate automatically. The next topic covers learning rate schedules — more advanced strategies for controlling exactly how the learning rate changes throughout training.

Leave a Comment

Your email address will not be published. Required fields are marked *