feat: add lr_scheduler

Author: chungongyu

profold2/command/trainer.py

@@ -23,7 +23,7 @@
 from profold2.data.utils import (
     embedding_get_labels, tensor_to_numpy, weights_from_file
 )
-from profold2.model import accelerator, FeatureBuilder, MetricDict, ReturnValues
+from profold2.model import accelerator, optim, FeatureBuilder, MetricDict, ReturnValues
 from profold2.model.utils import CheckpointManager
 from profold2.utils import exists
 

profold2/model/optim.py

@@ -0,0 +1,57 @@
+"""lr_scheduler wrapper
+"""
+from enum import Enum
+import math
+from typing import Optional, Union
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR
+
+from profold2.utils import default, exists
+
+
+class SchedulerType(Enum):
+  CONSTANT = 'constant'
+  COSINE = 'cosine'
+  LINEAR = 'linear'
+
+
+def get_scheduler(
+    name: Union[str, SchedulerType],
+    optimizer: Optimizer,
+    num_warmup_steps: Optional[int] = None,
+    num_training_steps: Optional[int] = None,
+    eta_min: float = 0.0,
+    last_epoch: int = -1,
+) -> LambdaLR:
+  name = SchedulerType(name)
+
+  if name == SchedulerType.CONSTANT:
+
+    def lr_lambda(current_step: int) -> float:
+      if exists(num_warmup_steps) and current_step < num_warmup_steps:
+        return current_step / max(1.0, num_warmup_steps)
+      return 1.0
+  elif name == SchedulerType.COSINE:
+
+    def lr_lambda(current_step: int) -> float:
+      if exists(num_warmup_steps) and current_step < num_warmup_steps:
+        return current_step / max(1.0, num_warmup_steps)
+      num_warmup_steps = default(num_warmup_steps, 0)
+      progress = (
+          (current_step - num_warmup_steps) / (num_training_steps - num_warmup_steps)
+      )
+      return 0.5  * (1.0 - eta_min) * (1.0 + math.cos(math.pi * progress)) + eta_min
+  elif name == SchedulerType.LINEAR:
+
+    def lr_lambda(current_step: int) -> float:
+      if exists(num_warmup_steps) and current_step < num_warmup_steps:
+        return current_step / max(1.0, num_warmup_steps)
+
+      num_warmup_steps = default(num_warmup_steps, 0)
+      progress = (
+          (num_training_steps - current_step) / (num_training_steps - num_warmup_steps)
+      )
+      return (1.0 - eta_min) * progress + eta_min
+
+  return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)