One-Cycle Learning Rate

Problem

Compute a one-cycle learning rate with linear warmup followed by cosine decay.

For 0 <= step < warmup_steps return min_lr + (base_lr - min_lr) * step / warmup_steps. For later steps let progress = (step - warmup_steps) / (total_steps - warmup_steps) and return min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * progress)). If total_steps is not positive or step >= total_steps return min_lr. When warmup_steps is not positive, decay starts from step 0.

Hint

Warmup ends at base_lr, then the cosine decay runs over the remaining steps.

Test cases

case 1

in: [0.1,0.001,10,100,0]

expected: 0.001

case 2

in: [0.1,0.001,10,100,10]

expected: 0.1

case 3

in: [0.1,0.001,10,100,55]

expected: 0.0505

case 4

in: [0.1,0.001,10,100,100]

expected: 0.001

case 5

in: [0.1,0.001,10,100,20]

expected: 0.09701478472890247

Starter codePython

def one_cycle_lr(base_lr, min_lr, warmup_steps, total_steps, step):
    # Your code here
    pass

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