tfm.optimization.CosineDecayWithOffset.base_lr_class

A LearningRateSchedule that uses a cosine decay with optional warmup.

tfm.optimization.CosineDecayWithOffset.base_lr_class(
    initial_learning_rate,
    decay_steps,
    alpha=0.0,
    name=None,
    warmup_target=None,
    warmup_steps=0
)

See Loshchilov & Hutter, ICLR2016, SGDR: Stochastic Gradient Descent with Warm Restarts.

For the idea of a linear warmup of our learning rate, see Goyal et al..

When we begin training a model, we often want an initial increase in our learning rate followed by a decay. If warmup_target is an int, this schedule applies a linear increase per optimizer step to our learning rate from initial_learning_rate to warmup_target for a duration of warmup_steps. Afterwards, it applies a cosine decay function taking our learning rate from warmup_target to alpha for a duration of decay_steps. If warmup_target is None we skip warmup and our decay will take our learning rate from initial_learning_rate to alpha. It requires a step value to compute the learning rate. You can just pass a TensorFlow variable that you increment at each training step.

The schedule is a 1-arg callable that produces a warmup followed by a decayed learning rate when passed the current optimizer step. This can be useful for changing the learning rate value across different invocations of optimizer functions.

Our warmup is computed as:

def warmup_learning_rate(step):
    completed_fraction = step / warmup_steps
    total_delta = target_warmup - initial_learning_rate
    return completed_fraction * total_delta

And our decay is computed as:

if warmup_target is None:
    initial_decay_lr = initial_learning_rate
else:
    initial_decay_lr = warmup_target

def decayed_learning_rate(step):
    step = min(step, decay_steps)
    cosine_decay = 0.5 * (1 + cos(pi * step / decay_steps))
    decayed = (1 - alpha) * cosine_decay + alpha
    return initial_decay_lr * decayed

Example usage without warmup:

decay_steps = 1000
initial_learning_rate = 0.1
lr_decayed_fn = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate, decay_steps)

Example usage with warmup:

decay_steps = 1000
initial_learning_rate = 0
warmup_steps = 1000
target_learning_rate = 0.1
lr_warmup_decayed_fn = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate, decay_steps, warmup_target=target_learning_rate,
    warmup_steps=warmup_steps
)

You can pass this schedule directly into a tf.keras.optimizers.Optimizer as the learning rate. The learning rate schedule is also serializable and deserializable using tf.keras.optimizers.schedules.serialize and tf.keras.optimizers.schedules.deserialize.

A 1-arg callable learning rate schedule that takes the current optimizer step and outputs the decayed learning rate, a scalar Tensor of the same type as initial_learning_rate.

initial_learning_rate A scalar float32 or float64 Tensor or a Python int. The initial learning rate.
decay_steps A scalar int32 or int64 Tensor or a Python int. Number of steps to decay over.
alpha A scalar float32 or float64 Tensor or a Python int. Minimum learning rate value for decay as a fraction of initial_learning_rate.
name String. Optional name of the operation. Defaults to 'CosineDecay'.
warmup_target None or a scalar float32 or float64 Tensor or a Python int. The target learning rate for our warmup phase. Will cast to the initial_learning_rate datatype. Setting to None will skip warmup and begins decay phase from initial_learning_rate. Otherwise scheduler will warmup from initial_learning_rate to warmup_target.
warmup_steps A scalar int32 or int64 Tensor or a Python int. Number of steps to warmup over.

Methods

from_config

@classmethod
from_config(
    config
)

Instantiates a LearningRateSchedule from its config.

Args
config Output of get_config().

Returns
A LearningRateSchedule instance.

get_config

get_config()

__call__

__call__(
    step
)

Call self as a function.