Fix the learning rate warmup schedule
Your task
An AI wrote a training loop with learning rate warmup, where the rate climbs from near zero to base_lr over the first steps and then holds. Its loop started at step 1, stopped one step early, and the rate never reached its peak. Your team writes the schedule out as a list, one rate per step, so it can be checked.
An assistant wrote warmup_lr_schedule(total_steps, warmup_steps, base_lr). It is in the editor. It has a bug. Find it and fix it.
Input and output:
total_steps: how many training steps run.warmup_steps: how many steps the warmup takes.base_lr: the peak learning rate.- Return the learning rate of each step, in order.
Rules the tests check:
- Return exactly
total_stepsrates, or an empty list iftotal_stepsis 0 or less. - Number the steps from 1. Step k's rate is
base_lrtimes the smaller of k divided bywarmup_steps, and 1. So the first step already trains, and the rate reachesbase_lrexactly at stepwarmup_steps. - If
warmup_stepsis 0 or less, every rate isbase_lr. - Answers within 0.000001 of the expected values pass.
Examples
warmup_lr_schedule(total_steps: int, warmup_steps: int, base_lr: float) -> list[float]Example 1
Input
warmup_lr_schedule(6, 4, 1)Output
[0.25, 0.5, 0.75, 1, 1, 1]Example 2
Input
warmup_lr_schedule(3, 5, 0.5)Output
[0.1, 0.2, 0.3]