Polynomial Decay Learning Rate

Problem

Compute a polynomial learning-rate decay toward end_lr:

lr = (base_lr - end_lr) * (1 - progress)^power + end_lr

where progress = step / total_steps clamped to [0, 1]. If total_steps is not positive, return end_lr.

Hint

power = 1 gives linear decay, power = 2 quadratic decay.

case 1

in: [0.1,0,100,1,0]

expected: 0.1

case 2

in: [0.1,0,100,1,50]

expected: 0.05

case 3

in: [0.1,0,100,2,50]

expected: 0.025

case 4

in: [0.1,0.01,100,0.5,100]

expected: 0.01

case 5

in: [0.1,0,100,1,150]

expected: 0

Starter codePython

def polynomial_decay_lr(base_lr, end_lr, total_steps, power, step):
    # Your code here
    pass

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