title writing your own callbacks s rick chao francois chollet date created 20190320 last modified 20230625 description complete guide to writing new keras callbacks accelerator gpu introduction a callback is a powerful tool to customize the behavior of a keras model during training evaluation or inference examples include keras callbacks tensorboard to visualize training progress and results with tensorboard or keras callbacks modelcheckpoint to periodically save your model during training in this guide you will learn what a keras callback is what it can do and how you can build your own we provide a few demos of simple callback applications to get you started setup keras callbacks overview all callbacks subclass the keras callbacks callback class and override a set of methods called at various stages of training testing and predicting callbacks are useful to get a view on internal states and statistics of the model during training you can pass a list of callbacks as the keyword argument callbacks to the following model methods keras model fit keras model evaluate keras model predict an overview of callback methods global methods ontraintestpredictbeginself logsnone called at the beginning of fitevaluatepredict ontraintestpredictendself logsnone called at the end of fitevaluatepredict batchlevel methods for trainingtestingpredicting ontraintestpredictbatchbeginself batch logsnone called right before processing a batch during trainingtestingpredicting ontraintestpredictbatchendself batch logsnone called at the end of trainingtestingpredicting a batch within this method logs is a dict containing the metrics results epochlevel methods training only onepochbeginself epoch logsnone called at the beginning of an epoch during training onepochendself epoch logsnone called at the end of an epoch during training a basic example let s take a look at a concrete example to get started let s import tensorflow and define a simple sequential keras model define the keras model to add callbacks to then load the mnist data for training and testing from keras datasets api load example mnist data and preprocess it limit the data to 1000 samples now define a simple custom callback that logs when fitevaluatepredict starts ends when each epoch starts ends when each training batch starts ends when each evaluation test batch starts ends when each inference prediction batch starts ends let s try it out usage of logs dict the logs dict contains the loss value and all the metrics at the end of a batch or epoch example includes the loss and mean absolute error usage of self model attribute in addition to receiving log information when one of their methods is called callbacks have access to the model associated with the current round of trainingevaluationinference self model here are a few of the things you can do with self model in a callback set self model stoptraining true to immediately interrupt training mutate hyperparameters of the optimizer available as self model optimizer such as self model optimizer learningrate save the model at period intervals record the output of model predict on a few test samples at the end of each epoch to use as a sanity check during training extract visualizations of intermediate features at the end of each epoch to monitor what the model is learning over time etc let s see this in action in a couple of examples examples of keras callback applications early stopping at minimum loss this first example shows the creation of a callback that stops training when the minimum of loss has been reached by setting the attribute self model stoptraining boolean optionally you can provide an argument patience to specify how many epochs we should wait before stopping after having reached a local minimum keras callbacks earlystopping provides a more complete and general implementation stop training when the loss is at its min i e the loss stops decreasing arguments patience number of epochs to wait after min has been hit after this number of no improvement training stops bestweights to store the weights at which the minimum loss occurs the number of epoch it has waited when loss is no longer minimum the epoch the training stops at initialize the best as infinity record the best weights if current results is better less learning rate scheduling in this example we show how a custom callback can be used to dynamically change the learning rate of the optimizer during the course of training see callbacks learningratescheduler for a more general implementations learning rate scheduler which sets the learning rate according to schedule arguments schedule a function that takes an epoch index integer indexed from 0 and current learning rate as inputs and returns a new learning rate as output float get the current learning rate from model s optimizer call schedule function to get the scheduled learning rate set the value back to the optimizer before this epoch starts epoch to start learning rate tuples helper function to retrieve the scheduled learning rate based on epoch if epoch lrschedule00 or epoch lrschedule10 return lr for i in rangelenlrschedule if epoch lrschedulei0 return lrschedulei1 return lr model getmodel model fit xtrain ytrain batchsize64 epochs15 verbose0 callbacks lossanderrorprintingcallback customlearningrateschedulerlrschedule builtin keras callbacks introduction a callback is a powerful tool to customize the behavior of a keras model during training evaluation or inference examples include keras callbacks tensorboard to visualize training progress and results with tensorboard or keras callbacks modelcheckpoint to periodically save your model during training in this guide you will learn what a keras callback is what it can do and how you can build your own we provide a few demos of simple callback applications to get you started setup keras callbacks overview all callbacks subclass the keras callbacks callback class and override a set of methods called at various stages of training testing and predicting callbacks are useful to get a view on internal states and statistics of the model during training you can pass a list of callbacks as the keyword argument callbacks to the following model methods keras model fit keras model evaluate keras model predict an overview of callback methods global methods on_ train test predict _begin self logs none called at the beginning of fit evaluate predict on_ train test predict _end self logs none called at the end of fit evaluate predict batch level methods for training testing predicting on_ train test predict _batch_begin self batch logs none called right before processing a batch during training testing predicting on_ train test predict _batch_end self batch logs none called at the end of training testing predicting a batch within this method logs is a dict containing the metrics results epoch level methods training only on_epoch_begin self epoch logs none called at the beginning of an epoch during training on_epoch_end self epoch logs none called at the end of an epoch during training a basic example let s take a look at a concrete example to get started let s import tensorflow and define a simple sequential keras model define the keras model to add callbacks to then load the mnist data for training and testing from keras datasets api load example mnist data and pre process it limit the data to 1000 samples now define a simple custom callback that logs when fit evaluate predict starts ends when each epoch starts ends when each training batch starts ends when each evaluation test batch starts ends when each inference prediction batch starts ends let s try it out usage of logs dict the logs dict contains the loss value and all the metrics at the end of a batch or epoch example includes the loss and mean absolute error usage of self model attribute in addition to receiving log information when one of their methods is called callbacks have access to the model associated with the current round of training evaluation inference self model here are a few of the things you can do with self model in a callback set self model stop_training true to immediately interrupt training mutate hyperparameters of the optimizer available as self model optimizer such as self model optimizer learning_rate save the model at period intervals record the output of model predict on a few test samples at the end of each epoch to use as a sanity check during training extract visualizations of intermediate features at the end of each epoch to monitor what the model is learning over time etc let s see this in action in a couple of examples examples of keras callback applications early stopping at minimum loss this first example shows the creation of a callback that stops training when the minimum of loss has been reached by setting the attribute self model stop_training boolean optionally you can provide an argument patience to specify how many epochs we should wait before stopping after having reached a local minimum keras callbacks earlystopping provides a more complete and general implementation stop training when the loss is at its min i e the loss stops decreasing arguments patience number of epochs to wait after min has been hit after this number of no improvement training stops best_weights to store the weights at which the minimum loss occurs the number of epoch it has waited when loss is no longer minimum the epoch the training stops at initialize the best as infinity record the best weights if current results is better less learning rate scheduling in this example we show how a custom callback can be used to dynamically change the learning rate of the optimizer during the course of training see callbacks learningratescheduler for a more general implementations learning rate scheduler which sets the learning rate according to schedule arguments schedule a function that takes an epoch index integer indexed from 0 and current learning rate as inputs and returns a new learning rate as output float get the current learning rate from model s optimizer call schedule function to get the scheduled learning rate set the value back to the optimizer before this epoch starts epoch to start learning rate tuples helper function to retrieve the scheduled learning rate based on epoch built in keras callbacks be sure to check out the existing keras callbacks by reading the api docs https keras io api callbacks applications include logging to csv saving the model visualizing metrics in tensorboard and a lot more
import numpy as np
import keras









def get_model():
    model = keras.Sequential()
    model.add(keras.layers.Dense(1))
    model.compile(
        optimizer=keras.optimizers.RMSprop(learning_rate=0.1),
        loss="mean_squared_error",
        metrics=["mean_absolute_error"],
    )
    return model





(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype("float32") / 255.0
x_test = x_test.reshape(-1, 784).astype("float32") / 255.0


x_train = x_train[:1000]
y_train = y_train[:1000]
x_test = x_test[:1000]
y_test = y_test[:1000]




class CustomCallback(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        keys = list(logs.keys())
        print("Starting training; got log keys: {}".format(keys))

    def on_train_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop training; got log keys: {}".format(keys))

    def on_epoch_begin(self, epoch, logs=None):
        keys = list(logs.keys())
        print(
            "Start epoch {} of training; got log keys: {}".format(epoch, keys)
        )

    def on_epoch_end(self, epoch, logs=None):
        keys = list(logs.keys())
        print("End epoch {} of training; got log keys: {}".format(epoch, keys))

    def on_test_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start testing; got log keys: {}".format(keys))

    def on_test_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop testing; got log keys: {}".format(keys))

    def on_predict_begin(self, logs=None):
        keys = list(logs.keys())
        print("Start predicting; got log keys: {}".format(keys))

    def on_predict_end(self, logs=None):
        keys = list(logs.keys())
        print("Stop predicting; got log keys: {}".format(keys))

    def on_train_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Training: start of batch {}; got log keys: {}".format(
                batch, keys
            )
        )

    def on_train_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Training: end of batch {}; got log keys: {}".format(batch, keys)
        )

    def on_test_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Evaluating: start of batch {}; got log keys: {}".format(
                batch, keys
            )
        )

    def on_test_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Evaluating: end of batch {}; got log keys: {}".format(
                batch, keys
            )
        )

    def on_predict_batch_begin(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Predicting: start of batch {}; got log keys: {}".format(
                batch, keys
            )
        )

    def on_predict_batch_end(self, batch, logs=None):
        keys = list(logs.keys())
        print(
            "...Predicting: end of batch {}; got log keys: {}".format(
                batch, keys
            )
        )




model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=128,
    epochs=1,
    verbose=0,
    validation_split=0.5,
    callbacks=[CustomCallback()],
)

res = model.evaluate(
    x_test, y_test, batch_size=128, verbose=0, callbacks=[CustomCallback()]
)

res = model.predict(x_test, batch_size=128, callbacks=[CustomCallback()])




class LossAndErrorPrintingCallback(keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        print(
            "Up to batch {}, the average loss is {:7.2f}.".format(
                batch, logs["loss"]
            )
        )

    def on_test_batch_end(self, batch, logs=None):
        print(
            "Up to batch {}, the average loss is {:7.2f}.".format(
                batch, logs["loss"]
            )
        )

    def on_epoch_end(self, epoch, logs=None):
        print(
            "The average loss for epoch {} is {:7.2f} "
            "and mean absolute error is {:7.2f}.".format(
                epoch, logs["loss"], logs["mean_absolute_error"]
            )
        )


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=128,
    epochs=2,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback()],
)

res = model.evaluate(
    x_test,
    y_test,
    batch_size=128,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback()],
)






class EarlyStoppingAtMinLoss(keras.callbacks.Callback):
    

    def __init__(self, patience=0):
        super().__init__()
        self.patience = patience
        
        self.best_weights = None

    def on_train_begin(self, logs=None):
        
        self.wait = 0
        
        self.stopped_epoch = 0
        
        self.best = np.Inf

    def on_epoch_end(self, epoch, logs=None):
        current = logs.get("loss")
        if np.less(current, self.best):
            self.best = current
            self.wait = 0
            
            self.best_weights = self.model.get_weights()
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.model.stop_training = True
                print("Restoring model weights from the end of the best epoch.")
                self.model.set_weights(self.best_weights)

    def on_train_end(self, logs=None):
        if self.stopped_epoch > 0:
            print(f"Epoch {self.stopped_epoch + 1}: early stopping")


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=30,
    verbose=0,
    callbacks=[LossAndErrorPrintingCallback(), EarlyStoppingAtMinLoss()],
)




class CustomLearningRateScheduler(keras.callbacks.Callback):
    

    def __init__(self, schedule):
        super().__init__()
        self.schedule = schedule

    def on_epoch_begin(self, epoch, logs=None):
        if not hasattr(self.model.optimizer, "learning_rate"):
            raise ValueError('Optimizer must have a "learning_rate" attribute.')
        
        lr = self.model.optimizer.learning_rate
        
        scheduled_lr = self.schedule(epoch, lr)
        
        self.model.optimizer.learning_rate = scheduled_lr
        print(
            f"\nEpoch {epoch}: Learning rate is {float(np.array(scheduled_lr))}."
        )


LR_SCHEDULE = [
    
    (3, 0.05),
    (6, 0.01),
    (9, 0.005),
    (12, 0.001),
]


def lr_schedule(epoch, lr):
    
    if epoch < LR_SCHEDULE[0][0] or epoch > LR_SCHEDULE[-1][0]:
        return lr
    for i in range(len(LR_SCHEDULE)):
        if epoch == LR_SCHEDULE[i][0]:
            return LR_SCHEDULE[i][1]
    return lr


model = get_model()
model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=15,
    verbose=0,
    callbacks=[
        LossAndErrorPrintingCallback(),
        CustomLearningRateScheduler(lr_schedule),
    ],
)