Coverage for cafl4ds/schedule.py: 100%

1"""Learning-rate schedules for the training loops. 

2 

3Phase-0 P0.1/P0.2 ran a **flat** learning rate over a single ~25-step pass. P0.2.1 found that 

4a genuine healthy SimSiam baseline (a RankMe that recovers and holds, rather than decaying 

5indistinguishably from a forced collapse) needs a *fair* training regime: multiple epochs with 

6a **warmup + cosine-decay** schedule. As the LR anneals, the intact model's effective rank 

7re-expands from the early-training dip — the healthy U-shape that separates it from the 

8collapsed arm. This module supplies that schedule as a plain, Hydra-instantiable factory. 

9""" 

10 

11from __future__ import annotations 

12 

13import math 

14 

15from torch import optim 

16from torch.optim.lr_scheduler import LambdaLR 

17 

18 

19def warmup_cosine_schedule( 

20 optimizer: optim.Optimizer, 

21 total_steps: int, 

22 warmup_frac: float = 0.1, 

23 min_lr_frac: float = 0.0, 

24) -> LambdaLR: 

25 """Linear warmup then cosine decay, expressed as a multiplier on the base LR. 

26 

27 The returned scheduler is stepped **once per optimizer step** (not per epoch). The 

28 multiplier ramps linearly from ``0`` to ``1`` over the first ``warmup_frac`` of 

29 ``total_steps``, then follows a half-cosine down to ``min_lr_frac`` at the final step. 

30 

31 Args: 

32 optimizer: The optimizer whose base LR(s) are scaled. 

33 total_steps: Total number of optimizer steps over the whole run 

34 (``epochs * batches_per_epoch``). 

35 warmup_frac: Fraction of ``total_steps`` spent warming up, in ``[0, 1)``. 

36 min_lr_frac: Final LR as a fraction of the base LR (``0`` decays fully to zero). 

37 

38 Returns: 

39 A :class:`~torch.optim.lr_scheduler.LambdaLR` implementing the schedule. 

40 

41 Raises: 

42 ValueError: If ``total_steps`` is not positive or ``warmup_frac`` is out of range. 

43 """ 

44 if total_steps <= 0: 

45 raise ValueError(f"total_steps must be positive, got {total_steps}.") 

46 if not 0.0 <= warmup_frac < 1.0: 

47 raise ValueError(f"warmup_frac must be in [0, 1), got {warmup_frac}.") 

48 warmup_steps = max(1, int(round(warmup_frac * total_steps))) 

49 

50 def lr_lambda(step: int) -> float: 

51 if step < warmup_steps: 

52 return (step + 1) / warmup_steps 

53 progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) 

54 cosine = 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) 

55 return min_lr_frac + (1.0 - min_lr_frac) * cosine 

56 

57 return LambdaLR(optimizer, lr_lambda)