Module 6 — Callbacks: checkpointing, early stopping, learning rate scheduling

A callback is an object that fit invokes at specific moments: the start and end of an epoch, of a batch, of training. It is how you act on a running training job without rewriting its loop. Three callbacks cover almost every need, and their defaults are rarely the right ones.

Save the best state, not the last

A training run of several hours that dies without a checkpoint is lost work. But the main point lies elsewhere: the last epoch's model is almost never the best one.

from tensorflow import keras

checkpoint = keras.callbacks.ModelCheckpoint(
    filepath="models/best.keras",
    monitor="val_loss",
    mode="min",
    save_best_only=True,
    save_weights_only=False,
    verbose=1,
)

save_best_only=True is the setting that matters. Without it, every epoch overwrites the file, and the last write is the final state — often overfitted. With it, the file holds the epoch where the monitored metric was best.

monitor and mode go together and contradict each other easily. Monitoring val_loss with mode="max" keeps the worst model, with no message to warn you. The rule: a loss is minimised, an accuracy is maximised. When in doubt, mode="auto" infers the direction from the metric name.

Stop when validation stops improving

EarlyStopping halts training when the monitored metric stagnates, which saves compute and limits overfitting.

stopping = keras.callbacks.EarlyStopping(
    monitor="val_loss",
    patience=10,
    min_delta=1e-4,
    restore_best_weights=True,
    verbose=1,
)

restore_best_weights=True is essential and defaults to False. Without it, training stops after patience stagnant epochs, and the model in memory is the last one, ten epochs past the best point. You detected the right moment to stop and kept the wrong model.

Patience is set according to how noisy your curves are. Too short, it cuts on an ordinary fluctuation when the loss was about to fall again. Too long, it serves no purpose. Five to fifteen epochs is a reasonable range, to be tuned after observing validation variability.

min_delta defines what counts as an improvement. Without it, a gain of $10^{-8}$ resets the patience counter and stopping never triggers.

Two callbacks that contradict each other

ModelCheckpoint and EarlyStopping must monitor the same metric. Monitoring val_loss for one and val_accuracy for the other gives you a stop governed by one quantity and a save governed by another: the file you keep then bears no relation to the moment training ended.

Reduce the learning rate at the right time

A fixed rate eventually prevents fine convergence: the steps stay too large to settle into the minimum. Two strategies exist.

The reactive one watches validation and cuts the rate when it stalls:

reduction = keras.callbacks.ReduceLROnPlateau(
    monitor="val_loss",
    factor=0.5,
    patience=5,
    min_lr=1e-6,
    verbose=1,
)

The scheduled one follows a curve decided in advance, independent of results:

schedule = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=1e-3,
    decay_steps=num_epochs * steps_per_epoch,
    alpha=0.01,
)
model.compile(optimizer=keras.optimizers.Adam(schedule), loss="mse")

Approach	Advantage	Limitation
reactive	adapts to what actually happens	reacts after the plateau, with lag
scheduled	reproducible, no patience hyperparameter	requires knowing the epoch count

Cosine decay has become the convention in large-model training, usually preceded by a warmup of a few hundred steps. That warmup prevents a full rate applied to freshly initialised weights from destroying information in the first few batches.

ReduceLROnPlateau patience must be lower than EarlyStopping patience

If early stopping is more impatient than the rate reduction, training ends before the first cut and the callback never fires. A ratio of two to three works well: patience 5 for the reduction, 12 for the stop.

Combining them, and knowing the order

model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=200,
    callbacks=[checkpoint, stopping, reduction,
               keras.callbacks.CSVLogger("log.csv")],
)

With this set of callbacks, epochs=200 is no longer a prediction but an upper bound: EarlyStopping decides when to finish. That is the right way to think about it, and it removes a hyperparameter from your list.

Callbacks run in list order at each epoch end. Putting ModelCheckpoint first guarantees the save happens before another callback interrupts training.

Writing your own callback

The interface is open, and a homemade callback fits in a few lines.

class StopAtThreshold(keras.callbacks.Callback):
    def __init__(self, threshold=0.98):
        super().__init__()
        self.threshold = threshold

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        if logs.get("val_accuracy", 0) >= self.threshold:
            print(f"\nThreshold reached at epoch {epoch + 1}, stopping.")
            self.model.stop_training = True

The logs dictionary holds the epoch's metrics, with the same names as in the history. self.model gives access to the full model, and self.model.stop_training = True is the very mechanism EarlyStopping uses.

The available hooks run from on_train_begin to on_predict_batch_end. Be careful with per-batch hooks, though: expensive work in on_train_batch_end executes thousands of times per epoch and can dominate training time.

Key takeaways

• ModelCheckpoint is only useful with save_best_only=True; otherwise the final file holds the last epoch, almost never the best.
• EarlyStopping requires restore_best_weights=True, which defaults to false: without it you detect the right stopping point and keep the wrong model.
• Rate reduction can be reactive with ReduceLROnPlateau or scheduled with cosine decay; in the first case its patience must stay below the early stopping patience.
• With these callbacks, epochs becomes an upper bound rather than a prediction, which removes a hyperparameter.

Next module: TensorBoard, for what the logs do not tell you.