fix lr_scheduler for BertAdam optimizer

alibaba · Jul 2, 2022 · 169bf0a · 169bf0a
1 parent bf17be3
commit 169bf0a
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/easynlp/core/optimizers.py b/easynlp/core/optimizers.py
@@ -495,7 +495,7 @@ def get_optimizer(
         optimizer_grouped_parameters.append({
             'params': [p],
             'weight_decay':
-            0.0 if any(nd in n for nd in no_decay) else weight_decay
+            0.0 if len(p.shape) == 1 or any(nd in n for nd in no_decay) else weight_decay
         })
     if optimizer_type == 'BertAdam':
         optimizer = BertAdam(

diff --git a/easynlp/core/trainer.py b/easynlp/core/trainer.py
@@ -107,6 +107,7 @@ def set_model_and_optimizer(self, model, args):
         # Build Optimizer
         self._optimizer, self._lr_scheduler = get_optimizer(
             optimizer_type=self.optimizer_type,
+            schedule=args.lr_scheduler,
             learning_rate=args.learning_rate,
             warmup_proportion=args.warmup_proportion,
             max_grad_norm=self.max_grad_norm,

diff --git a/easynlp/utils/arguments.py b/easynlp/utils/arguments.py
@@ -296,6 +296,16 @@ def _add_easynlp_args(parser: argparse.ArgumentParser):
                        ],
                        help='name of the optimizer')
 
+    group.add_argument('--lr_scheduler',
+                       '--schedule',
+                       default='warmup_linear',
+                       type=str,
+                       choices=[
+                           'warmup_linear', 'warmup_cosine',
+                           'warmup_constant', 'none',
+                       ],
+                       help='name of the learning rate scheduler')
+
     group.add_argument(
         '--warmup_proportion',
         default=0.1,