Neural Mastery
← Practice
Linear Learning-Rate Warmup Scheduleeasy

Linear Learning-Rate Warmup Schedule

easy ML Systems & MLOps 15 min

NumPy allowed · plain Python earns 10 bonus points

Implement linear learning-rate warmup, a real, near-universal training stabilization technique for large model training.

Task

Implement warmup_lr(step, warmup_steps, base_lr). If step >= warmup_steps (or warmup_steps == 0), return base_lr. Otherwise return base_lr * step / warmup_steps (linear ramp from 0).

Function Signature

warmup_lr(step: int, warmup_steps: int, base_lr: float) -> float

Examples

Example 1: midway through warmup
Input: step = 50 warmup_steps = 100 base_lr = 0.001
Output: 0.0005
Example 2: at step 0
Input: step = 0 warmup_steps = 100 base_lr = 0.001
Output: 0
Example 3: past warmup, full lr
Input: step = 200 warmup_steps = 100 base_lr = 0.001
Output: 0.001

Constraints

  • step >= 0
  • warmup_steps >= 0
Python3Saved ✓
def warmup_lr(step, warmup_steps, base_lr):
# Your implementation here
pass

Case 1: midway through warmup
Input: step = 50 warmup_steps = 100 base_lr = 0.001
Expected: 0.0005
Case 2: at step 0
Input: step = 0 warmup_steps = 100 base_lr = 0.001
Expected: 0
Case 3: past warmup, full lr
Input: step = 200 warmup_steps = 100 base_lr = 0.001
Expected: 0.001