From 169bf0a706985a06fc2598960fdafc8c3be64944 Mon Sep 17 00:00:00 2001 From: Ling-Hao Han Date: Sat, 2 Jul 2022 16:52:15 +0800 Subject: [PATCH] fix lr_scheduler for BertAdam optimizer --- easynlp/core/optimizers.py | 2 +- easynlp/core/trainer.py | 1 + easynlp/utils/arguments.py | 10 ++++++++++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/easynlp/core/optimizers.py b/easynlp/core/optimizers.py index 77bcb33..dec0dc5 100644 --- 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 index e31c0b5..5f16b68 100644 --- 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 index ccf245c..05281dc 100644 --- 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,