add point finetune: lr schedule from dynamic to pre-computed

Author: Mark-ZhouWX

official/cv/segment-anything/configs/sa1b_point_finetune.yaml

@@ -23,7 +23,7 @@ optimizer:
   group_param:
 
   lr_scheduler:
-    type: segment_anything.optim.scheduler.SAMDynamicDecayLR
+    type: segment_anything.optim.scheduler.sam_dynamic_decay_lr
     learning_rate: 8e-6
     warmup_steps: 250
     decay_steps: [ 60000, 86666 ]

official/cv/segment-anything/segment_anything/optim/optimizer.py

@@ -10,6 +10,8 @@
 def create_optimizer(
     params,
     args,
+    step_per_epoch,
+    epoch_size
 ):
     r"""Creates optimizer by name.
 
@@ -25,15 +27,16 @@ def create_optimizer(
     Returns:
         Optimizer object
     """
-    optimizer = OPTIMIZER_REGISTRY.instantiate(**args, params=params)
+    optimizer = OPTIMIZER_REGISTRY.instantiate(**args, params=params,
+                                               step_per_epoch=step_per_epoch, epoch_size=epoch_size)
     return optimizer
 
 
 @OPTIMIZER_REGISTRY.registry_module()
 class AdamW(nn.optim.Adam):
-    def __init__(self, params: List, lr_scheduler, group_param, **kwargs):
+    def __init__(self, params: List, lr_scheduler, group_param, step_per_epoch, epoch_size, **kwargs):
         if group_param is None:
             group_param = dict()
         params = create_group_param(params, **group_param)
-        lr_scheduler_inst = create_lr_scheduler(lr_scheduler)
+        lr_scheduler_inst = create_lr_scheduler(lr_scheduler, step_per_epoch=step_per_epoch, epoch_size=epoch_size)
         super().__init__(params, lr_scheduler_inst, **kwargs)

official/cv/segment-anything/segment_anything/optim/scheduler.py

@@ -7,15 +7,42 @@
 from segment_anything.utils.registry import LR_SCHEDULER_REGISTRY
 
 
-def create_lr_scheduler(args: Dict):
+def create_lr_scheduler(args, step_per_epoch, epoch_size):
     """
     instantiate learning rate scheduler class
     """
+    if args.type.endswith('sam_dynamic_decay_lr'):
+        return sam_dynamic_decay_lr(learning_rate=args.learning_rate,
+                                    warmup_steps=args.warmup_steps,
+                                    decay_steps=args.decay_steps,
+                                    decay_factor=args.decay_factor,
+                                    step_per_epoch=step_per_epoch,
+                                    epoch_size=epoch_size,
+                                    )
     scheduler = LR_SCHEDULER_REGISTRY.instantiate(**args)
     return scheduler
 
+def sam_dynamic_decay_lr(learning_rate, warmup_steps, decay_steps, decay_factor, step_per_epoch, epoch_size):
+    def lr_factor(step):
+        if step < warmup_steps:
+            return step / float(warmup_steps)
+        elif step < decay_steps[0]:
+            return 1.0
+        elif step < decay_steps[1]:
+            return 1.0 / decay_factor
+        else:
+            return 1.0 / (decay_factor**2)
+    total_step = step_per_epoch * epoch_size
+    lr_list = []
+    for i in range(total_step):
+        step = i + 1
+        lr = learning_rate * lr_factor(step)
+        lr_list.append(lr)
+
+    return lr_list
+
 
-@LR_SCHEDULER_REGISTRY.registry_module()
+# @LR_SCHEDULER_REGISTRY.registry_module()
 class SAMDynamicDecayLR(LearningRateSchedule):
     def __init__(self,
                  learning_rate: float,

official/cv/segment-anything/segment_anything/utils/callbacks.py

@@ -130,7 +130,7 @@ def on_train_step_end(self, run_context: RunContext):
         self.accumulate_loss += loss
 
         if cur_step % self.log_interval == 0:
-            lr = cb_params.network.optimizer.learning_rate(cur_step)
+            lr = cb_params.network.optimizer.learning_rate.learning_rate[cur_step]
             smooth_loss = self.accumulate_loss / self.log_interval
 
             step_cost = time.time() - self.step_start_time

official/cv/segment-anything/train.py

@@ -37,7 +37,9 @@ def main(args) -> None:
     network = amp.auto_mixed_precision(network, args.get('amp_level', 'O0'))
 
     # Step3: create optimizer, including learning rate scheduler and group parameter settings
-    optimizer = create_optimizer(params=network.trainable_params(), args=args.optimizer)
+    optimizer = create_optimizer(params=network.trainable_params(), args=args.optimizer,
+                                 step_per_epoch=train_dataloader.get_dataset_size(),
+                                 epoch_size=args.train_loader.epoch_size)
 
     # Step4: wrap model and optimizer for training
     with_loss_model = NetWithLossWrapper(network, loss_fn=loss_fn,