update LRScheduler for rnnlm (#4984)

* update lrscheduler

* minor fix

* add pre-commit

PaddleNLP/examples/language_model/rnnlm/model.py

@@ -87,13 +87,6 @@ class CrossEntropyLossForLm(nn.Layer):
 
 class UpdateModel(paddle.callbacks.Callback):
     # This callback reset model hidden states and update learning rate before each epoch begins 
-    def __init__(self, base_lr, lr_decay, epoch_start_decay):
-        self.base_lr = base_lr
-        self.lr_decay = lr_decay
-        self.epoch_start_decay = epoch_start_decay
-
     def on_epoch_begin(self, epoch=None, logs=None):
         self.model.network.reset_states()
-        new_lr = self.base_lr * (self.lr_decay
-                                 **max(epoch + 1 - self.epoch_start_decay, 0.0))
-        self.model._optimizer.set_lr(new_lr)
+        

PaddleNLP/examples/language_model/rnnlm/train.py

@@ -13,8 +13,9 @@ paddle.seed(102)
 
 
 def create_data_loader(batch_size, num_steps, data_path):
-    train_ds, valid_ds, test_ds = PTBDataset.get_datasets(
-        [batch_size] * 3, [num_steps] * 3, ['train', 'eval', 'test'])
+    train_ds = PTBDataset(batch_size, num_steps, 'train')
+    valid_ds = PTBDataset(batch_size, num_steps, 'eval')
+    test_ds = PTBDataset(batch_size, num_steps, 'test')
 
     train_loader = DataLoader(train_ds, return_list=True, batch_size=None)
     valid_loader = DataLoader(valid_ds, return_list=True, batch_size=None)
@@ -40,15 +41,15 @@ def train(args):
     gloabl_norm_clip = paddle.nn.ClipGradByGlobalNorm(args.max_grad_norm)
     cross_entropy = CrossEntropyLossForLm()
     ppl_metric = Perplexity()
-    callback = UpdateModel(
-        base_lr=args.base_lr,
-        lr_decay=args.lr_decay,
-        epoch_start_decay=args.epoch_start_decay)
-
+    callback = UpdateModel()
+    scheduler = paddle.callbacks.LRScheduler(by_step=False, by_epoch=True)
     model = paddle.Model(network)
-    # FIXME(yuanxiaopeng): Use scheduler instead of callback
-    #scheduler = paddle.optimizer.lr.LambdaDecay(learning_rate=args.base_lr, lr_lambda=lambda x: args.lr_decay**max(x + 1 - args.epoch_start_decay, 0.0), verbose=True)
-    optimizer = paddle.optimizer.SGD(learning_rate=args.base_lr,
+
+    learning_rate = paddle.optimizer.lr.LambdaDecay(
+        learning_rate=args.base_lr,
+        lr_lambda=lambda x: args.lr_decay**max(x + 1 - args.epoch_start_decay, 0.0),
+        verbose=True)
+    optimizer = paddle.optimizer.SGD(learning_rate=learning_rate,
                                      parameters=model.parameters(),
                                      grad_clip=gloabl_norm_clip)
 
@@ -62,7 +63,7 @@ def train(args):
               eval_data=valid_loader,
               epochs=args.max_epoch,
               shuffle=False,
-              callbacks=[callback],
+              callbacks=[callback, scheduler],
               log_freq=max(1, len(train_loader) // 10))
 
     model.save(path='checkpoint/test')  # save for training