diff --git a/skillopt/optimizer/scheduler.py b/skillopt/optimizer/scheduler.py index 416d3bb..50d1072 100644 --- a/skillopt/optimizer/scheduler.py +++ b/skillopt/optimizer/scheduler.py @@ -63,16 +63,15 @@ class ConstantScheduler(LRScheduler): class LinearScheduler(LRScheduler): """Linear decay from ``max_lr`` to ``min_lr`` over ``total_steps``. - The scheduler guarantees that ``step()`` returns ``max_lr`` on the first - call and ``min_lr`` on the ``total_steps``-th call. Intermediate steps - are linearly interpolated and rounded to the nearest integer. + For multi-step runs, the first call evaluates the first decay interval + (``t = 1 / total_steps``) and the ``total_steps``-th call returns + ``min_lr``. Values are rounded to the nearest integer. """ def _compute_lr(self, step: int) -> int: if self.total_steps <= 1: return self.max_lr - s = min(step, self.total_steps) - t = (s - 1) / (self.total_steps - 1) + t = min(step, self.total_steps) / self.total_steps lr = self.max_lr + (self.min_lr - self.max_lr) * t return max(self.min_lr, round(lr)) @@ -80,16 +79,15 @@ class LinearScheduler(LRScheduler): class CosineScheduler(LRScheduler): """Cosine annealing from ``max_lr`` to ``min_lr`` over ``total_steps``. - The scheduler guarantees that ``step()`` returns ``max_lr`` on the first - call and ``min_lr`` on the ``total_steps``-th call. Intermediate steps - follow a half-cosine curve that starts and ends flat. + For multi-step runs, the first call evaluates the first decay interval + (``t = 1 / total_steps``) and the ``total_steps``-th call returns + ``min_lr``. Intermediate steps follow a half-cosine curve. """ def _compute_lr(self, step: int) -> int: if self.total_steps <= 1: return self.max_lr - s = min(step, self.total_steps) - t = (s - 1) / (self.total_steps - 1) + t = min(step, self.total_steps) / self.total_steps lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + math.cos(math.pi * t)) return max(self.min_lr, round(lr)) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 40bf3ec..29a5a7c 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -9,8 +9,9 @@ math — making it an ideal target for precise unit tests. Scheduler contract ------------------ -Decay schedulers (Linear, Cosine) guarantee: -- First ``step()`` returns ``max_lr``. +For multi-step runs, decay schedulers (Linear, Cosine) guarantee: +- First ``step()`` evaluates the first decay interval at + ``t = 1 / total_steps`` (rounding can still produce ``max_lr``). - ``total_steps``-th ``step()`` returns ``min_lr``. - Intermediate values are monotonically non-increasing and clamped to ``[min_lr, max_lr]``. @@ -81,9 +82,9 @@ class TestConstantScheduler: class TestLinearScheduler: """LinearScheduler — linear decay from max_lr to min_lr.""" - def test_first_step_returns_max_lr(self) -> None: + def test_first_step_uses_first_decay_interval(self) -> None: s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10) - assert s.step() == 10 + assert s.step() == 9 def test_last_step_returns_min_lr(self) -> None: s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10) @@ -134,14 +135,14 @@ class TestLinearScheduler: def test_known_decay_sequence(self) -> None: """Linear decay max_lr=10, min_lr=2, total_steps=4. - t = (step-1)/(total_steps-1) = 0, 1/3, 2/3, 1 + t = step/total_steps = 1/4, 1/2, 3/4, 1 lr = 10 + (2-10)*t = 10 - 8t - t=0: lr=10, t=1/3: lr≈7.33→7, t=2/3: lr≈4.67→5, t=1: lr=2 + t=1/4: lr=8, t=1/2: lr=6, t=3/4: lr=4, t=1: lr=2 """ s = LinearScheduler(max_lr=10, min_lr=2, total_steps=4) - assert s.step() == 10 - assert s.step() == 7 - assert s.step() == 5 + assert s.step() == 8 + assert s.step() == 6 + assert s.step() == 4 assert s.step() == 2 def test_step_state_dict_resume_consistent(self) -> None: @@ -167,9 +168,9 @@ class TestLinearScheduler: class TestCosineScheduler: """CosineScheduler — cosine annealing from max_lr to min_lr.""" - def test_first_step_returns_max_lr(self) -> None: - s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10) - assert s.step() == 10 + def test_first_step_uses_first_decay_interval(self) -> None: + s = CosineScheduler(max_lr=100, min_lr=0, total_steps=4) + assert s.step() == 85 def test_last_step_returns_min_lr(self) -> None: s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10) @@ -218,8 +219,8 @@ class TestCosineScheduler: def test_midpoint_close_to_mean(self) -> None: """At the half-way neighbourhood, cosine is close to (max+min)/2. - total_steps=100, step=50 → t=49/99≈0.495. - cos(0.495π)≈0, lr ≈ (20+2)/2 = 11. + total_steps=100, step=50 → t=50/100=0.5. + cos(0.5π)=0, lr = (20+2)/2 = 11. """ s = CosineScheduler(max_lr=20, min_lr=2, total_steps=100) for _ in range(49): @@ -243,13 +244,17 @@ class TestCosineScheduler: assert s.step() == 5 def test_early_steps_near_max(self) -> None: - """Cosine annealing stays near max_lr early on (cos(0)=1).""" + """Cosine annealing stays near max_lr early on.""" s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100) - # step 1: t=0, cos(0)=1 → lr=100 + # step 1: t=0.01, cos(0.01π)≈0.9995 → lr≈99.98 → 100 assert s.step() == 100 - # step 2: t=1/99≈0.01, cos≈0.9995 → lr≈99.97 → 100 + # step 2: t=0.02, cos(0.02π)≈0.9980 → lr≈99.90 → 100 assert s.step() == 100 + def test_known_decay_sequence(self) -> None: + s = CosineScheduler(max_lr=10, min_lr=2, total_steps=4) + assert [s.step() for _ in range(4)] == [9, 6, 3, 2] + def test_late_steps_near_min(self) -> None: """Cosine annealing flattens near min_lr at the end (cos(π)=-1).""" s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)