Warmup and Cosine Decay

Write the learning rate schedule that most large models train with: a linear warmup, then a cosine decay down to a floor.

For a 0-based step number step:

• Warmup, when step < warmup_steps: the rate climbs in a straight line, base_lr * (step + 1) / warmup_steps. Using step + 1 means the first step does not have a rate of zero.
• Decay, when warmup_steps <= step < total_steps: let progress = (step - warmup_steps) / (total_steps - warmup_steps). The rate is min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * progress)).
• After the end, when step >= total_steps: the rate stays at min_lr.

Function signature

def warmup_cosine_lr(step, base_lr, warmup_steps, total_steps, min_lr=0.0):

Example

Input

[warmup_cosine_lr(s, 0.001, 10, 110) for s in (0, 9, 10, 60, 110)]

Output

[0.0001, 0.001, 0.001, 0.0005, 0.0]

Reasoning

Step 0 is the first warmup step, so the rate is 0.001 * 1 / 10. Step 9 ends the warmup at the full rate. Step 10 starts the decay with progress = 0, so the cosine term is 1 and the rate is still 0.001. Step 60 is half way through the 100 decay steps, where cos(pi / 2) = 0, so the rate is half. Step 110 is past the end, so the rate is min_lr.

Constraints and notes

• warmup_steps can be 0. Then the schedule is pure cosine decay.
• You can assume warmup_steps < total_steps.
• Return a plain float.