Warmup-Stable-Decay Learning Rate

Problem

Compute the learning rate of a warmup-stable-decay (WSD) schedule.

During the first warmup_steps the rate rises linearly from 0 to base_lr, then stays at base_lr for stable_steps, then decays linearly to 0 over decay_steps, and is 0 afterwards. Steps before 0 are treated as 0.

Hint

The three phases are contiguous at step values warmup, warmup + stable, and the end.

Test cases

case 1

in: [0.1,10,20,30,0]

expected: 0

case 2

in: [0.1,10,20,30,5]

expected: 0.05

case 3

in: [0.1,10,20,30,10]

expected: 0.1

case 4

in: [0.1,10,20,30,30]

expected: 0.1

case 5

in: [0.1,10,20,30,45]

expected: 0.05

case 6

in: [0.1,10,20,30,60]

expected: 0

Starter codePython

def wsd_lr(base_lr, warmup_steps, stable_steps, decay_steps, step):
    # Your code here
    pass

6 test cases · run your code against them in the editor