add files

funasr_local/schedulers/noam_lr.py

@@ -0,0 +1,65 @@
+"""Noam learning rate scheduler module."""
+from typing import Union
+import warnings
+
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+from typeguard import check_argument_types
+
+from funasr_local.schedulers.abs_scheduler import AbsBatchStepScheduler
+
+
+class NoamLR(_LRScheduler, AbsBatchStepScheduler):
+    """The LR scheduler proposed by Noam
+
+    Ref:
+        "Attention Is All You Need", https://arxiv.org/pdf/1706.03762.pdf
+
+    FIXME(kamo): PyTorch doesn't provide _LRScheduler as public class,
+     thus the behaviour isn't guaranteed at forward PyTorch version.
+
+    NOTE(kamo): The "model_size" in original implementation is derived from
+     the model, but in this implementation, this parameter is a constant value.
+     You need to change it if the model is changed.
+
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        model_size: Union[int, float] = 320,
+        warmup_steps: Union[int, float] = 25000,
+        last_epoch: int = -1,
+    ):
+        assert check_argument_types()
+        self.model_size = model_size
+        self.warmup_steps = warmup_steps
+
+        lr = list(optimizer.param_groups)[0]["lr"]
+        new_lr = self.lr_for_WarmupLR(lr)
+        warnings.warn(
+            f"NoamLR is deprecated. "
+            f"Use WarmupLR(warmup_steps={warmup_steps}) with Optimizer(lr={new_lr})",
+        )
+
+        # __init__() must be invoked before setting field
+        # because step() is also invoked in __init__()
+        super().__init__(optimizer, last_epoch)
+
+    def lr_for_WarmupLR(self, lr: float) -> float:
+        return lr / self.model_size**0.5 / self.warmup_steps**0.5
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(model_size={self.model_size}, "
+            f"warmup_steps={self.warmup_steps})"
+        )
+
+    def get_lr(self):
+        step_num = self.last_epoch + 1
+        return [
+            lr
+            * self.model_size**-0.5
+            * min(step_num**-0.5, step_num * self.warmup_steps**-1.5)
+            for lr in self.base_lrs
+        ]