Add source: tmt/training/scheduler.py

tmt/training/scheduler.py

@@ -0,0 +1,28 @@
+"""
+scheduler.py — Cosine annealing with linear warmup.
+"""
+from __future__ import annotations
+
+import math
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+
+def cosine_warmup_scheduler(
+    optimizer: Optimizer,
+    warmup_steps: int,
+    total_steps: int,
+    min_lr_ratio: float = 0.1,
+) -> LambdaLR:
+    """
+    Linear warmup for `warmup_steps`, then cosine decay to `min_lr_ratio * lr`.
+    """
+    def lr_lambda(step: int) -> float:
+        if step < warmup_steps:
+            return step / max(warmup_steps, 1)
+        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
+        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine
+
+    return LambdaLR(optimizer, lr_lambda)