| import math | |
| from mla.schedule import lr_schedule | |
| def test_warmup_linear(): | |
| assert abs(lr_schedule(0, 1.0, 10, 110) - 0.1) < 1e-12 | |
| assert abs(lr_schedule(4, 1.0, 10, 110) - 0.5) < 1e-12 | |
| assert abs(lr_schedule(9, 1.0, 10, 110) - 1.0) < 1e-12 | |
| def test_peak_at_warmup_boundary(): | |
| assert abs(lr_schedule(10, 1.0, 10, 110) - 1.0) < 1e-12 | |
| def test_cosine_midpoint(): | |
| assert abs(lr_schedule(60, 1.0, 10, 110, min_lr=0.0) - 0.5) < 1e-12 | |
| def test_cosine_quarter(): | |
| assert abs(lr_schedule(35, 1.0, 10, 110) - 0.5 * (1.0 + math.cos(math.pi * 0.25))) < 1e-12 | |
| def test_decays_to_min(): | |
| assert abs(lr_schedule(110, 1.0, 10, 110) - 0.0) < 1e-12 | |
| assert abs(lr_schedule(500, 1.0, 10, 110) - 0.0) < 1e-12 | |
| def test_min_lr_floor(): | |
| assert abs(lr_schedule(110, 1.0, 10, 110, min_lr=0.1) - 0.1) < 1e-12 | |
| assert abs(lr_schedule(60, 1.0, 10, 110, min_lr=0.2) - (0.2 + 0.8 * 0.5)) < 1e-12 | |
| def test_monotonic_non_increasing_after_warmup(): | |
| xs = [lr_schedule(s, 1.0, 10, 110) for s in range(10, 111)] | |
| assert all(xs[i] >= xs[i + 1] - 1e-15 for i in range(len(xs) - 1)) | |