Keras Custom Callbacks

You've trained models with EarlyStopping and ModelCheckpoint before, but what happens when you need to pause training to refill your GPU's memory, or send a Slack alert when validation loss plateaus? The built-in Keras callbacks cover the common cases, but they can't know about your training loop's unique demands. That's the pain: you hit a wall where the training behavior you need simply isn't bundled with tf.keras. The solution is to create custom callbacks in Keras — a skill that turns a generic training script into a purpose-built, production-grade pipeline. In this Applied AI engineering lesson, you'll learn the exact mechanism Keras uses to hook into every stage of training, and you'll walk away able to write callbacks that monitor, mutate, and even halt training on your own terms.

The problem this lesson solves

Training a neural network is rarely a fire-and-forget operation. In real projects, you need to:

• React to metrics as they emerge — e.g., reduce the learning rate when validation loss stops improving.
• Save the best model only when it truly is the best (not just every epoch).
• Log custom information to your own dashboard, CSV, or messaging system.
• Implement business rules — like stopping training early if a custom fairness metric crosses a threshold, or if GPU memory usage exceeds a safety limit.
• Interact with the outside world — send a notification, trigger a data pipeline step, or adjust a hyperparameter mid-run.

The built-in callbacks (EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard) cover generic scenarios, but they are black boxes. You can set parameters, but you can't inject your own logic at precise moments. When you hit this ceiling, you have two choices: hack together a loop that manually checks conditions between epochs (ugly and error-prone), or write a custom callback that integrates seamlessly with the Keras training engine. The latter is cleaner, more testable, and far more maintainable — and that's exactly what this lesson teaches.

Core concept / mental model

Think of a Keras callback as a set of hooks — functions that Keras promises to call at specific moments during training. Picture a robot that trains your model. At every checkpoint (start of training, end of each epoch, end of each batch, etc.), the robot stops, looks at its list of observers, and says: "Hey, does anyone want to do something right now?" Your custom callback is one such observer, raising its hand at the moments you care about.

Formally, every callback inherits from tf.keras.callbacks.Callback. When you pass a list of callbacks to model.fit(), Keras invokes the corresponding methods on each callback at the right time. The most important methods are:

logs is a dictionary that Keras populates with the current metrics — e.g., logs['loss'], logs['accuracy'], logs['val_loss']. The key mental shift: you are not writing code that runs inside the training loop; you are writing code that Keras calls at the boundaries. This separation keeps your architecture clean and lets you mix multiple callbacks without entangling their effects.

How it works step by step

To create a custom callback, follow this systematic path:

1. Subclass tf.keras.callbacks.Callback — your class inherits all the methods, and you override the ones you need.
2. Define __init__ — store configuration values (like a threshold, a file path, or a log name).
3. Override the lifecycle methods — choose the hook points that matter for your goal.
4. Access the model and logs — inside the callback, self.model gives you the current model; logs gives you the metrics.
5. Use self.model.stop_training = True — this is the magic flag that halts training from inside a callback.
6. Pass an instance to model.fit() — along with any built-in callbacks, in the callbacks list.

Pro tip: The logs dictionary is only populated with metrics that are actually computed. If you want val_accuracy to appear, make sure you pass a validation set to fit().

A simple example

Here's the minimal custom callback that prints the training loss after every epoch:

import tensorflow as tf

class LossLogger(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        loss = logs.get("loss")
        print(f"Epoch {epoch + 1}: loss = {loss:.4f}")

# Train a mini model to see the callback in action
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(2,))])
model.compile(optimizer="sgd", loss="mse")

# Dummy data
import numpy as np
x = np.random.rand(100, 2)
y = np.random.rand(100, 1)

model.fit(x, y, epochs=3, callbacks=[LossLogger()])

Expected output (abbreviated):

Epoch 1: loss = 0.1234
Epoch 2: loss = 0.1187
Epoch 3: loss = 0.1123

Hands-on walkthrough

Now let's build something genuinely useful: a custom callback that saves the best model based on a custom metric and stops training if a metric plateaus — but with your own twist. You'll also see how to modify the learning rate dynamically.

Example 1: Custom ModelCheckpoint with a Twist

Keras's built-in ModelCheckpoint saves based on a monitored quantity, but let's say you want to keep the last N best models, not just the single best. Here's a custom callback:

import tensorflow as tf
import numpy as np
import os

class KeepLastNBest(tf.keras.callbacks.Callback):
    def __init__(self, monitor="val_loss", n=3, save_path="checkpoints"):
        super().__init__()
        self.monitor = monitor
        self.n = n
        self.save_path = save_path
        self.history = []  # list of (metric_value, filename)
        os.makedirs(save_path, exist_ok=True)

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is None:
            return
        # Save the model with epoch number
        filename = f"{self.save_path}/model_epoch_{epoch + 1}.keras"
        self.model.save(filename)
        self.history.append((current, filename))
        # Sort by metric (ascending for loss, descending for accuracy)
        self.history.sort(key=lambda x: x[0], reverse=("acc" in self.monitor))
        # Remove extras
        while len(self.history) > self.n:
            _, old_file = self.history.pop()
            if os.path.exists(old_file):
                os.remove(old_file)
        print(f"Saved {filename} (val_loss: {current:.4f})")

# Usage
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(2,))])
model.compile(optimizer="adam", loss="mse")

x = np.random.rand(100, 2)
y = np.random.rand(100, 1)
val_x = np.random.rand(20, 2)
val_y = np.random.rand(20, 1)

model.fit(x, y, validation_data=(val_x, val_y), epochs=5,
          callbacks=[KeepLastNBest(n=2)])

Expected behavior: After training, only the 2 best models remain in the checkpoints/ folder.

Example 2: Learning Rate Scheduler with Custom Logic

You can implement your own learning rate scheduler that reduces LR after a plateau, but with a minimum limit and a cooldown period:

class CustomLRScheduler(tf.keras.callbacks.Callback):
    def __init__(self, factor=0.5, patience=2, min_lr=1e-6, cooldown=1):
        super().__init__()
        self.factor = factor
        self.patience = patience
        self.min_lr = min_lr
        self.cooldown = cooldown
        self.wait = 0
        self.cooldown_counter = 0
        self.best = float("inf")

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        val_loss = logs.get("val_loss")
        if val_loss is None:
            return

        if self.cooldown_counter > 0:
            self.cooldown_counter -= 1
            return

        if val_loss < self.best:
            self.best = val_loss
            self.wait = 0
        else:
            self.wait += 1
            if self.wait >= self.patience:
                lr = float(tf.keras.backend.get_value(self.model.optimizer.lr))
                new_lr = max(lr * self.factor, self.min_lr)
                tf.keras.backend.set_value(self.model.optimizer.lr, new_lr)
                print(f"Reducing LR to {new_lr:.2e}")
                self.wait = 0
                self.cooldown_counter = self.cooldown

# Usage
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(2,))])
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.01), loss="mse")

model.fit(x, y, validation_data=(val_x, val_y), epochs=10,
          callbacks=[CustomLRScheduler()])

Example 3: Early Stopping with a Custom Condition

Let's stop training not based on patience, but when a custom metric (like a business KPI) crosses a threshold:

class ThresholdStopper(tf.keras.callbacks.Callback):
    def __init__(self, threshold=0.95, monitor="accuracy"):
        super().__init__()
        self.threshold = threshold
        self.monitor = monitor

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is not None and current >= self.threshold:
            print(f"Reached {self.monitor} = {current:.4f} >= {self.threshold}. Stopping.")
            self.model.stop_training = True

# Usage with a small classification model
inputs = tf.keras.Input(shape=(4,))
x = tf.keras.layers.Dense(8, activation="relu")(inputs)
outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
model = tf.keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

x_cls = np.random.rand(100, 4)
y_cls = np.random.randint(0, 2, size=(100,))
model.fit(x_cls, y_cls, epochs=50, callbacks=[ThresholdStopper(threshold=0.8)])

Compare options / when to choose what

When should you write a custom callback instead of using built-ins? Here's a comparison table:

General rule of thumb: If a built-in does exactly what you need, use it. But if you find yourself combining multiple built-ins in a hacky way or trying to make one behave like another, a custom callback is the cleaner path.

Variations

• TensorFlow 1.x style: Older tf.compat.v1 had a Callback class with similar methods. If you're working on legacy code, the patterns are almost identical, but you won't have the logs dict populated the same way.
• Keras 2.x vs 3.x: The Callback base class is the same API. In Keras 3 (multi-backend), tf.keras is one implementation; you can also use keras directly with PyTorch or JAX backends, and the callback API remains the same.
• Custom training loops: If you use tf.GradientTape instead of model.fit(), you have full control and may not need callbacks — but you lose the convenience of metric tracking. Use callbacks when you want a standard training loop with injected logic.

Troubleshooting & edge cases

• logs is empty or missing keys: This happens when the metric isn't computed. For validation metrics, you must pass .fit(validation_data=...). For custom metrics, ensure they're included in model.compile(metrics=[...]).
• Fix: Always check logs.get("val_loss") or logs["loss"] as a fallback.
• self.model is None in __init__: The model isn't attached until before on_train_begin. Don't access self.model in __init__; do it in on_train_begin or later.
• Setting stop_training doesn't stop: Make sure you set self.model.stop_training = True in a method that actually runs during training (e.g., on_epoch_end). Setting it in __init__ has no effect.
• Modifying the model's loss or state inside a callback can cause unexpected behavior: If you change weights or add layers, do it in on_epoch_begin or on_train_begin, and be aware that changing the architecture after compilation may break optimizers.
• Memory leak when saving models every batch: Use on_epoch_end for heavy operations unless you know what you're doing. Batch-level callbacks run thousands of times — optimize accordingly.
• Concurrency issues with shared state: If you use multiple callbacks that share state (e.g., a counter), keep the state inside each callback instance to avoid side effects.

What you learned & what's next

You now have the core skill of creating custom callbacks in Keras. You learned the mental model of hooks, the key lifecycle methods (on_epoch_begin, on_epoch_end, on_batch_end, etc.), how to access logs and self.model, and how to stop training or modify the learning rate. You also compared when to use built-ins vs. custom, and you walked through three concrete examples: a rolling-best model saver, a custom LR scheduler, and a threshold-based early stopper. This lesson directly fulfills the objectives: you can explain the core idea and you've completed practical exercises.

The next lesson in this track will likely cover advanced training loops — perhaps custom metrics or callbacks that interact with tf.data pipelines. With custom callbacks, you're now ready to debug, monitor, and control long training runs in production.