scheduler

Provides learning rate scheduling functions and classes.

This module contains implementations of linear and cosine learning rate schedules with optional warmup periods. These can be used to dynamically adjust learning rates during training of machine learning models.

Typical usage example:

schedule = CosineSchedule(max_learning_rate=6e-4, min_learning_rate=0,

total_steps=5000, warmup_steps=100)

learning_rate = schedule(current_step)

class CosineSchedule(max_learning_rate, min_learning_rate, total_steps, warmup_steps=0)

Bases: object

A class to implement a cosine decay learning rate schedule with warmup.

Parameters:

• max_learning_rate (float)

• min_learning_rate (float)

• total_steps (int)

• warmup_steps (int)

max_learning_rate

Maximum learning rate.

min_learning_rate

Minimum learning rate.

total_steps

Total number of steps in the training process.

warmup_steps

Number of warmup steps.

n_steps

Number of steps after warmup.

coef

Coefficient used in the cosine decay calculation.

step(step)

Calculates the learning rate for a given step.

Parameters:

step (int) – Current step in the training process.

Returns:

The calculated learning rate for the current step.

Return type:

float

linear_schedule(step, max_learning_rate, min_learning_rate, warmup_steps, total_steps)

Calculates the learning rate using a linear decay schedule with warmup.

Parameters:

• step (int) – Current step in the training process.

• max_learning_rate (float) – Maximum learning rate.

• min_learning_rate (float) – Minimum learning rate.

• warmup_steps (int) – Number of warmup steps.

• total_steps (int) – Total number of steps in the training process.

Returns:

The calculated learning rate for the current step.

Raises:

ValueError – If warmup_steps is greater than total_steps.

Return type:

float