Merge pull request #384 from bigict/optim

feat: add lr_scheduler

Author: chungongyu

profold2/command/trainer.py

@@ -25,7 +25,7 @@
 )
 from profold2.model import accelerator, optim, FeatureBuilder, MetricDict, ReturnValues
 from profold2.model.utils import CheckpointManager
-from profold2.utils import exists
+from profold2.utils import default, exists
 
 from profold2.command import worker
 
@@ -260,11 +260,11 @@ def model_params_groups(optim_options):
             break
       return params
 
-    optim = Adam(
+    optimizer = Adam(
         model_params_groups(args.model_params_optim_option), lr=args.learning_rate
     )
   else:
-    optim = Adam(model.parameters(), lr=args.learning_rate)
+    optimizer = Adam(model.parameters(), lr=args.learning_rate)
 
   # tensorboard
   writer = SummaryWriter(os.path.join(args.prefix, 'runs', 'eval')
@@ -310,12 +310,21 @@ def writer_add_scalars(writer, loss, it, prefix=''):
         os.path.join(args.prefix, 'checkpoints'),
         max_to_keep=args.checkpoint_max_to_keep,
         model=model,
-        optimizer=optim
+        optimizer=optimizer
     )
     global_step = checkpoint_manager.restore_or_initialize() + 1
     logging.info('checkpoint_manager.global_step: %d', global_step)
     model.train()
 
+  scheduler = optim.get_scheduler(
+      args.lr_scheduler,
+      optimizer,
+      num_warmup_steps=args.lr_scheduler_warmup_steps,
+      num_training_steps=default(args.lr_scheduler_training_steps, args.num_batches),
+      eta_min=args.lr_scheduler_eta_min,
+      last_global_step=global_step
+  )
+
   # .. note:: When a model is trained on ``M`` nodes with ``batch=N``, the
   #     gradient will be ``M`` times smaller when compared to the same model
   #     trained on a single node with ``batch=M*N`` if the loss is summed (NOT
@@ -331,9 +340,11 @@ def writer_add_scalars(writer, loss, it, prefix=''):
                  1) / (args.gradient_accumulate_every or 1.0)
 
   def _step(data_loader, it, writer, stage='train', batch_callback=None):
-    optim.zero_grad(set_to_none=True)
+    optimizer.zero_grad(set_to_none=True)
 
-    logging.debug('_step it: %d, loss_scaler: %f', it, loss_scaler)
+    logging.debug(
+        '_step it: %d, loss_scaler: %f, lr: %s', it, loss_scaler, scheduler.get_lr()
+    )
 
     running_loss = MetricDict()
     for jt in range(args.gradient_accumulate_every):
@@ -378,9 +389,10 @@ def _step(data_loader, it, writer, stage='train', batch_callback=None):
       writer_add_scalars(writer, v, it, prefix=f'Loss/{stage}@{k}')
       # writer.add_scalar(f'Loss/train@{k}', v, it)
 
-    # optim.step()
-    grad_scaler.step(optim)
+    # optimizer.step()
+    grad_scaler.step(optimizer)
     grad_scaler.update()
+    scheduler.step()
 
   def batch_seq_only(batch):
     batch = copy.copy(batch)
@@ -687,6 +699,31 @@ def add_arguments(parser):  # pylint: disable=redefined-outer-name
   parser.add_argument(
       '-l', '--learning_rate', type=float, default='1e-3', help='learning rate.'
   )
+  parser.add_argument(
+      '--lr_scheduler',
+      type=str,
+      default=optim.SchedulerType.CONSTANT.value,
+      choices=[m.value for m in optim.SchedulerType],
+      help='lr scheduler.'
+  )
+  parser.add_argument(
+      '--lr_scheduler_warmup_steps',
+      type=float,
+      default=None,
+      help='num of warmup steps for lr scheduler.'
+  )
+  parser.add_argument(
+      '--lr_scheduler_training_steps',
+      type=float,
+      default=None,
+      help='num of training steps for applying lr scheduler.'
+  )
+  parser.add_argument(
+      '--lr_scheduler_eta_min',
+      type=float,
+      default=0.0,
+      help='eta_min for applying lr scheduler.'
+  )
 
   parser.add_argument(
       '--model_features',

profold2/model/optim.py

@@ -1,6 +1,7 @@
 """lr_scheduler wrapper
 """
 from enum import Enum
+import functools
 import math
 from typing import Optional, Union
 
@@ -22,36 +23,51 @@ def get_scheduler(
     num_warmup_steps: Optional[int] = None,
     num_training_steps: Optional[int] = None,
     eta_min: float = 0.0,
-    last_epoch: int = -1,
+    last_global_step: int = 0,
 ) -> LambdaLR:
   name = SchedulerType(name)
 
   if name == SchedulerType.CONSTANT:
 
-    def lr_lambda(current_step: int) -> float:
+    def lr_lambda(
+        current_step: int, num_warmup_steps: Optional[int] = None
+    ) -> float:
+      current_step = current_step + last_global_step
       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:
+    def lr_lambda(
+        current_step: int, num_warmup_steps: Optional[int] = None
+    ) -> float:
+      current_step = current_step + last_global_step
       if exists(num_warmup_steps) and current_step < num_warmup_steps:
         return current_step / max(1.0, num_warmup_steps)
+      elif current_step > num_training_steps:
+        return eta_min
       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:
+    def lr_lambda(
+        current_step: int, num_warmup_steps: Optional[int] = None
+    ) -> float:
+      current_step = current_step + last_global_step
       if exists(num_warmup_steps) and current_step < num_warmup_steps:
         return current_step / max(1.0, num_warmup_steps)
+      elif current_step > num_training_steps:
+        return eta_min
 
       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)
+  return LambdaLR(
+      optimizer, functools.partial(lr_lambda, num_warmup_steps=num_warmup_steps)
+  )