Examples of callbacks use

Choice-Learn is compatible with the module tf.keras.callbacks. It allows the use of the different callbacks during the .fit() such as early-stopping or reduce-lr-on-plateau.\ In this notebook we showcase a few examples: - Early-Stopping

Otherwise the official documentation: - for callbacks - for custom callbacks

Early-Stopping

Early-Stopping allows the model to stop before reaching the given amount of epochs, when the monitored loss does not improve anymore.\ Here is an example.

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import sys
sys.path.append("../../")

import matplotlib.pyplot as plt
import tensorflow as tf
tf.config.run_functions_eagerly(True)

from choice_learn.datasets import load_swissmetro
from choice_learn.models import MultiNomialLogit

dataset = load_swissmetro()

callbacks = [
    tf.keras.callbacks.EarlyStopping(monitor="train_loss", patience=20, min_delta=0, mode="min")
]
model = MultiNomialLogit(epochs=1_000, optimizer="adam", callbacks=callbacks, lr=1e-3, batch_size=-1)
hist = model.fit(dataset)

plt.plot(hist["train_loss"])
plt.xlabel("Epochs")
plt.ylabel("Train Loss")
plt.show()

Reduce Learning Rate on Plateau

Reduce Learning Rate on Plateau allows the model to reduce the learning rate value, when the monitored loss does not improve anymore.\ Here is an example.

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import sys
sys.path.append("../../")

import matplotlib.pyplot as plt
import tensorflow as tf
tf.config.run_functions_eagerly(True)

from choice_learn.datasets import load_swissmetro
from choice_learn.models import MultiNomialLogit

dataset = load_swissmetro()

callbacks = [
    tf.keras.callbacks.ReduceLROnPlateau(monitor="train_loss", patience=20, mode="min")
]
model = MultiNomialLogit(epochs=1_000, optimizer="adam", callbacks=callbacks, lr=1e-1, batch_size=-1)

print("Starting model training, with learning rate:", model.optimizer.learning_rate)
hist = model.fit(dataset, verbose=1)
print("Ending model training, with learning rate:", model.optimizer.learning_rate)

plt.plot(hist["train_loss"])
plt.xlabel("Epochs")
plt.ylabel("Train Loss")
plt.show()

Customizing Callbacks

Let's show an example of a custom callback that stores the weights values at each epoch.

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
import sys
sys.path.append("../../")

import matplotlib.pyplot as plt
import tensorflow as tf
tf.config.run_functions_eagerly(True)

class ParamtersStore(tf.keras.callbacks.Callback):
    def __init__(self):
        super().__init__()
        self.params_store = {}

    def on_epoch_end(self, epoch, logs=None):
        # Store the parameters at the end of each epoch
        for params in self.model._trainable_weights:
            self.params_store.setdefault(params.name[:-2], {})[epoch] = params.numpy()

from choice_learn.datasets import load_swissmetro
from choice_learn.models import MultiNomialLogit

dataset = load_swissmetro()

callbacks = [
    ParamtersStore()
]
model = MultiNomialLogit(epochs=1_000, lr=1e-3, optimizer="adam", callbacks=callbacks, batch_size=-1)

hist = model.fit(dataset, verbose=1)

plt.figure(figsize=(16, 4))
plt.subplot(1, 3, 1)
plt.plot([callbacks[0].params_store["Weights_items_features"][epoch][1] for epoch in range(1_000)])
plt.title("Weight Items Features 2")
plt.xlabel("Epoch")
plt.ylabel("Weight Value")

plt.subplot(1, 3, 2)
plt.plot([callbacks[0].params_store["Weights_items_features"][epoch][3] for epoch in range(1_000)])
plt.title("Weight Items Features 4")
plt.xlabel("Epoch")
plt.ylabel("Weight Value")

plt.subplot(1, 3, 3)
plt.plot(hist["train_loss"], c="red")
plt.title("Train Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")

plt.tight_layout()
plt.show()