stouputils.data_science.models.keras_utils.callbacks.warmup_scheduler module#

class WarmupScheduler(warmup_epochs: int, initial_lr: float, target_lr: float)[source]#

Bases: Callback

Keras Callback for learning rate warmup.

Sources: - Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour: https://arxiv.org/abs/1706.02677 - Attention Is All You Need: https://arxiv.org/abs/1706.03762

This callback implements a learning rate warmup strategy where the learning rate gradually increases from an initial value to a target value over a specified number of epochs. This helps stabilize training in the early stages.

The learning rate increases linearly from the initial value to the target value over the warmup period, and then remains at the target value.

warmup_epochs: int#

Number of epochs for warmup.

initial_lr: float#

Starting learning rate for warmup.

target_lr: float#

Target learning rate after warmup.

model: Model#

Model to apply the warmup scheduler to.

on_epoch_begin(epoch: int, logs: dict[str, Any] | None = None) → None[source]#

Adjust learning rate at the beginning of each epoch during warmup.

Parameters:

  • epoch (int) – Current epoch index.

  • logs (dict | None) – Training logs.