Step LR scheduler

Background

A learning-rate scheduler changes the learning rate as training proceeds. StepLR is the simplest: it holds the rate constant, then multiplies it by a factor $\gamma$ every fixed number of epochs. Large early steps make fast progress; the periodic decay then lets the model settle into a minimum.

Problem statement

Implement a StepLRScheduler class with:

• __init__(self, initial_lr, step_size, gamma)
• get_lr(self, epoch) returning the learning rate at a given (0-indexed) epoch:

$$\text{lr}(e) = \text{lr}_0 \cdot \gamma^{\left\lfloor e / s \right\rfloor}$$

where $s$ is step_size. Round the result to 4 decimals.

Input

• initial_lr — float, the starting learning rate $\text{lr}_0$.
• step_size — int, the decay period (rate drops every step_size epochs).
• gamma — float, the multiplicative decay factor (e.g. 0.5 halves the rate).
• epoch — int, the current epoch (0-indexed) passed to get_lr.

Output

get_lr(epoch) returns a float learning rate, rounded to 4 decimals.

Examples

Example 1

Input:  StepLRScheduler(initial_lr=0.1, step_size=5, gamma=0.5)
        get_lr(0), get_lr(4), get_lr(5), get_lr(9), get_lr(10)
Output: 0.1, 0.1, 0.05, 0.05, 0.025

Explanation: the rate stays 0.1 for epochs 0–4 ($\lfloor e/5\rfloor=0$), drops to $0.1\cdot 0.5 = 0.05$ at epoch 5, and to $0.1\cdot 0.5^2 = 0.025$ at epoch 10.

Constraints

• The number of decays is the floor of epoch / step_size (integer division).
• The decay is multiplicative: $\gamma$ raised to the number of decays.
• Round the returned learning rate to 4 decimals.

Notes

• StepLR is a step function of the epoch: flat plateaus separated by sudden drops at multiples of step_size.
• The smooth alternative is ExponentialLR, which decays every epoch by $\gamma$ rather than in discrete steps.