add files

funasr_local/schedulers/abs_scheduler.py

@@ -0,0 +1,84 @@
+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):
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        pass
+
+
+# If you need to define custom scheduler, please inherit these classes
+class AbsBatchStepScheduler(AbsScheduler):
+    @abstractmethod
+    def step(self, epoch: int = None):
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        pass
+
+
+class AbsEpochStepScheduler(AbsScheduler):
+    @abstractmethod
+    def step(self, epoch: int = None):
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        pass
+
+
+class AbsValEpochStepScheduler(AbsEpochStepScheduler):
+    @abstractmethod
+    def step(self, val, epoch: int = None):
+        pass
+
+    @abstractmethod
+    def state_dict(self):
+        pass
+
+    @abstractmethod
+    def load_state_dict(self, state):
+        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)