Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/schedulers/__init__.py

@@ -0,0 +1,25 @@
+import torch
+import torch.multiprocessing
+import torch.nn
+import torch.optim
+
+from funasr.schedulers.noam_lr import NoamLR
+from funasr.schedulers.tri_stage_scheduler import TriStageLR
+from funasr.schedulers.warmup_lr import WarmupLR
+from funasr.schedulers.lambdalr_cus import CustomLambdaLR
+
+scheduler_classes = dict(
+    ReduceLROnPlateau=torch.optim.lr_scheduler.ReduceLROnPlateau,
+    lambdalr=torch.optim.lr_scheduler.LambdaLR,
+    steplr=torch.optim.lr_scheduler.StepLR,
+    multisteplr=torch.optim.lr_scheduler.MultiStepLR,
+    exponentiallr=torch.optim.lr_scheduler.ExponentialLR,
+    CosineAnnealingLR=torch.optim.lr_scheduler.CosineAnnealingLR,
+    noamlr=NoamLR,
+    warmuplr=WarmupLR,
+    tri_stage=TriStageLR,
+    cycliclr=torch.optim.lr_scheduler.CyclicLR,
+    onecyclelr=torch.optim.lr_scheduler.OneCycleLR,
+    CosineAnnealingWarmRestarts=torch.optim.lr_scheduler.CosineAnnealingWarmRestarts,
+    custom_lambdalr=CustomLambdaLR,
+)

funasr/schedulers/abs_scheduler.py

@@ -0,0 +1,129 @@
+from abc import ABC
+from abc import abstractmethod
+
+import torch.optim.lr_scheduler as L
+
+
+class AbsScheduler(ABC):
+    @abstractmethod
+    def step(self, epoch: int = None):
+        """Step.
+        
+            Args:
+                epoch: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        """State dict."""
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        """Load state dict.
+        
+            Args:
+                state: TODO.
+            """
+        pass
+
+
+# If you need to define custom scheduler, please inherit these classes
+class AbsBatchStepScheduler(AbsScheduler):
+    @abstractmethod
+    def step(self, epoch: int = None):
+        """Step.
+        
+            Args:
+                epoch: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        """State dict."""
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        """Load state dict.
+        
+            Args:
+                state: TODO.
+            """
+        pass
+
+
+class AbsEpochStepScheduler(AbsScheduler):
+    @abstractmethod
+    def step(self, epoch: int = None):
+        """Step.
+        
+            Args:
+                epoch: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        """State dict."""
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        """Load state dict.
+        
+            Args:
+                state: TODO.
+            """
+        pass
+
+
+class AbsValEpochStepScheduler(AbsEpochStepScheduler):
+    @abstractmethod
+    def step(self, val, epoch: int = None):
+        """Step.
+        
+            Args:
+                val: TODO.
+                epoch: TODO.
+            """
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        """State dict."""
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        """Load state dict.
+        
+            Args:
+                state: TODO.
+            """
+        pass
+
+
+# Create alias type to check the type
+# Note(kamo): Currently PyTorch doesn't provide the base class
+# to judge these classes.
+AbsValEpochStepScheduler.register(L.ReduceLROnPlateau)
+for s in [
+    L.ReduceLROnPlateau,
+    L.LambdaLR,
+    L.StepLR,
+    L.MultiStepLR,
+    L.MultiStepLR,
+    L.ExponentialLR,
+    L.CosineAnnealingLR,
+]:
+    AbsEpochStepScheduler.register(s)
+
+AbsBatchStepScheduler.register(L.CyclicLR)
+for s in [
+    L.OneCycleLR,
+    L.CosineAnnealingWarmRestarts,
+]:
+    AbsBatchStepScheduler.register(s)

funasr/schedulers/lambdalr_cus.py

@@ -0,0 +1,51 @@
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+# class CustomLambdaLR(_LRScheduler):
+#     def __init__(self, optimizer, warmup_steps, last_epoch=-1):
+#         self.warmup_steps = warmup_steps
+#         super().__init__(optimizer, last_epoch)
+#
+#     def get_lr(self):
+#         if self.last_epoch < self.warmup_steps:
+#             return [
+#                 base_lr * min(self.last_epoch / self.warmup_steps, 1) for base_lr in self.base_lrs
+#             ]
+#         else:
+#             return [base_lr for base_lr in self.base_lrs]
+
+
+class CustomLambdaLR(_LRScheduler):
+    def __init__(
+        self,
+        optimizer,
+        warmup_steps: int = 25000,
+        total_steps: int = 500000,
+        last_epoch=-1,
+        verbose=False,
+    ):
+        """Initialize CustomLambdaLR.
+        
+            Args:
+                optimizer: TODO.
+                warmup_steps: TODO.
+                total_steps: TODO.
+                last_epoch: TODO.
+                verbose: TODO.
+            """
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        super().__init__(optimizer, last_epoch, verbose)
+
+    def get_lr(self):
+
+        """Get lr."""
+        step = self.last_epoch + 1
+        if step < self.warmup_steps:
+            lr_scale = step / self.warmup_steps
+        else:
+            lr_scale = max(
+                0.0, 1 - (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+            )
+        return [base_lr * lr_scale for base_lr in self.base_lrs]

funasr/schedulers/noam_lr.py

@@ -0,0 +1,77 @@
+"""Noam learning rate scheduler module."""
+
+from typing import Union
+import warnings
+
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+from funasr.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,
+    ):
+        """Initialize NoamLR.
+        
+            Args:
+                optimizer: TODO.
+                model_size: Size/dimension parameter.
+                warmup_steps: TODO.
+                last_epoch: TODO.
+            """
+        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:
+        """Lr for warmuplr.
+        
+            Args:
+                lr: TODO.
+            """
+        return lr / self.model_size**0.5 / self.warmup_steps**0.5
+
+    def __repr__(self):
+        """Internal: repr  ."""
+        return (
+            f"{self.__class__.__name__}(model_size={self.model_size}, "
+            f"warmup_steps={self.warmup_steps})"
+        )
+
+    def get_lr(self):
+        """Get lr."""
+        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
+        ]

funasr/schedulers/tri_stage_scheduler.py

@@ -0,0 +1,124 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Optional, List
+
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler
+
+
+class TriStageLR(_LRScheduler, AbsBatchStepScheduler):
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        last_epoch: int = -1,
+        phase_ratio: Optional[List[float]] = None,
+        init_lr_scale: float = 0.01,
+        final_lr_scale: float = 0.01,
+    ):
+        """Initialize TriStageLR.
+        
+            Args:
+                optimizer: TODO.
+                last_epoch: TODO.
+                phase_ratio: TODO.
+                init_lr_scale: TODO.
+                final_lr_scale: TODO.
+            """
+        self.optimizer = optimizer
+        self.last_epoch = last_epoch
+        self.phase_ratio = phase_ratio
+        self.init_lr_scale = init_lr_scale
+        self.final_lr_scale = final_lr_scale
+        self.optimizer_lr = self.optimizer.defaults["lr"]
+
+    def init_tri_stage_scheudler(self, max_update):
+        """Init tri stage scheudler.
+        
+            Args:
+                max_update: TODO.
+            """
+        self.max_update = max_update
+        self.peak_lr = self.optimizer_lr
+        self.init_lr = self.init_lr_scale * self.optimizer_lr
+        self.final_lr = self.final_lr_scale * self.optimizer_lr
+
+        assert self.max_update > 0
+        assert sum(self.phase_ratio) == 1, "phase ratios must add up to 1"
+        assert len(self.phase_ratio) == 3
+        self.warmup_steps = int(self.max_update * self.phase_ratio[0])
+        self.hold_steps = int(self.max_update * self.phase_ratio[1])
+        self.decay_steps = int(self.max_update * self.phase_ratio[2])
+
+        self.warmup_rate = (
+            (self.peak_lr - self.init_lr) / self.warmup_steps if self.warmup_steps != 0 else 0
+        )
+        self.decay_factor = -math.log(self.final_lr_scale) / self.decay_steps
+
+        # initial learning rate
+        self.lr = self.init_lr
+
+        # __init__() must be invoked before setting field
+        # because step() is also invoked in __init__()
+        self.set_optimizer_lr(self.lr)
+        super().__init__(self.optimizer, self.last_epoch)
+
+    def _decide_stage(self, update_step):
+        """
+        return stage, and the corresponding steps within the current stage
+        """
+        if update_step < self.warmup_steps:
+            # warmup state
+            return 0, update_step
+
+        offset = self.warmup_steps
+
+        if update_step < offset + self.hold_steps:
+            # hold stage
+            return 1, update_step - offset
+
+        offset += self.hold_steps
+
+        if update_step <= offset + self.decay_steps:
+            # decay stage
+            return 2, update_step - offset
+
+        offset += self.decay_steps
+
+        # still here ? constant lr stage
+        return 3, update_step - offset
+
+    def step_update(self, num_updates):
+        """Update the learning rate after each update."""
+        stage, steps_in_stage = self._decide_stage(num_updates)
+        if stage == 0:
+            self.lr = self.init_lr + self.warmup_rate * steps_in_stage
+        elif stage == 1:
+            self.lr = self.peak_lr
+        elif stage == 2:
+            self.lr = self.peak_lr * math.exp(-self.decay_factor * steps_in_stage)
+        elif stage == 3:
+            self.lr = self.final_lr
+        else:
+            raise ValueError("Undefined stage")
+        self.set_optimizer_lr(self.lr)
+
+    def set_optimizer_lr(self, lr):
+        """Set optimizer lr.
+        
+            Args:
+                lr: TODO.
+            """
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+
+    def get_lr(self):
+        """Get lr."""
+        step_num = self.last_epoch + 1
+        self.step_update(step_num)
+        return [self.lr]

funasr/schedulers/warmup_lr.py

@@ -0,0 +1,56 @@
+"""Warm up learning rate scheduler module."""
+
+from typing import Union
+
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+
+from funasr.schedulers.abs_scheduler import AbsBatchStepScheduler
+
+
+class WarmupLR(_LRScheduler, AbsBatchStepScheduler):
+    """The WarmupLR scheduler
+
+    This scheduler is almost same as NoamLR Scheduler except for following difference:
+
+    NoamLR:
+        lr = optimizer.lr * model_size ** -0.5
+             * min(step ** -0.5, step * warmup_step ** -1.5)
+    WarmupLR:
+        lr = optimizer.lr * warmup_step ** 0.5
+             * min(step ** -0.5, step * warmup_step ** -1.5)
+
+    Note that the maximum lr equals to optimizer.lr in this scheduler.
+
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_steps: Union[int, float] = 25000,
+        last_epoch: int = -1,
+    ):
+        """Initialize WarmupLR.
+        
+            Args:
+                optimizer: TODO.
+                warmup_steps: TODO.
+                last_epoch: TODO.
+            """
+        self.warmup_steps = warmup_steps
+
+        # __init__() must be invoked before setting field
+        # because step() is also invoked in __init__()
+        super().__init__(optimizer, last_epoch)
+
+    def __repr__(self):
+        """Internal: repr  ."""
+        return f"{self.__class__.__name__}(warmup_steps={self.warmup_steps})"
+
+    def get_lr(self):
+        """Get lr."""
+        step_num = self.last_epoch + 1
+        return [
+            lr * self.warmup_steps**0.5 * min(step_num**-0.5, step_num * self.warmup_steps**-1.5)
+            for lr in self.base_lrs
+        ]