How to create a custom callback in Keras

Model training often needs project-specific logging or state capture without moving that behavior into a layer or loss function. A Keras callback receives lifecycle events from fit(), which keeps those side effects separate from the model definition.

Custom callbacks subclass keras.callbacks.Callback and override only the lifecycle methods they need. on_epoch_end() receives a zero-based epoch index and a logs dictionary containing the loss and compiled metrics for that epoch.

A small regression model is enough to exercise the callback without downloading a dataset. The retained epoch list and metric-key snapshots prove that the callback ran after every requested epoch and kept the observed training state.

1. Create custom_callback.py with the imports and the EpochLogRecorder callback class.

   custom_callback.py

   import keras
   import numpy as np
    
    
   class EpochLogRecorder(keras.callbacks.Callback):
       def __init__(self):
           super().__init__()
           self.epochs = []
           self.log_keys = []
    
       def on_train_begin(self, logs=None):
           self.epochs.clear()
           self.log_keys.clear()
    
       def on_epoch_end(self, epoch, logs=None):
           keys = sorted((logs or {}).keys())
           self.epochs.append(epoch + 1)
           self.log_keys.append(keys)
           print(f"epoch={epoch + 1} log_keys={','.join(keys)}")

   on_train_begin() resets callback-owned state for each fit() call, while on_epoch_end() records the epoch number and metric names supplied by Keras.

2. Append the deterministic training data and compiled regression model below the callback class.

   keras.utils.set_random_seed(7)
    
   x_train = np.arange(0, 8, dtype="float32").reshape(-1, 1)
   y_train = (2 * x_train) + 1
    
   model = keras.Sequential(
       [
           keras.layers.Input(shape=(1,)),
           keras.layers.Dense(1),
       ]
   )
   model.compile(
       optimizer=keras.optimizers.SGD(learning_rate=0.05),
       loss="mse",
       metrics=["mae"],
   )

3. Append the callback instance and fit() call below the model configuration.

   recorder = EpochLogRecorder()
   history = model.fit(
       x_train,
       y_train,
       batch_size=4,
       epochs=3,
       verbose=0,
       callbacks=[recorder],
   )

   The callbacks list registers each object for the lifecycle points implemented by its class.

4. Append the retained-state checks below the fit() call.

   assert recorder.epochs == [1, 2, 3]
   assert recorder.log_keys == [["loss", "mae"]] * 3
    
   print("recorded epochs:", recorder.epochs)
   print("recorded log keys:", recorder.log_keys)
   print("history keys:", sorted(history.history.keys()))

5. Run custom_callback.py to exercise the callback during three training epochs.

   $ python3 custom_callback.py
   epoch=1 log_keys=loss,mae
   epoch=2 log_keys=loss,mae
   epoch=3 log_keys=loss,mae
   recorded epochs: [1, 2, 3]
   recorded log keys: [['loss', 'mae'], ['loss', 'mae'], ['loss', 'mae']]
   history keys: ['loss', 'mae']

   One callback line per epoch plus the retained lists confirm that fit() invoked on_epoch_end() three times with both compiled training quantities.