StepLR Learning Rate Scheduler

Problem

Implement the StepLR learning rate scheduler. The learning rate is decayed by a factor gamma every step_size epochs.

Solution

def step_lr(initial_lr: float, step_size: int, gamma: float, epoch: int) -> float:
    return initial_lr * (gamma ** (epoch // step_size))

Explanation

1. At each epoch, compute how many complete step_size intervals have passed: epoch // step_size.
2. Multiply the initial learning rate by gamma raised to that power.
3. For example, with initial_lr=0.1, step_size=30, gamma=0.1: at epoch 0-29 the LR is 0.1, at epoch 30-59 it is 0.01, and so on.

Complexity

• Time: O(1)
• Space: O(1)