Minor update for transformer-xl (PaddlePaddle#192)

* minor update * delete useless break
wangxicoding · Mar 26, 2021 · 2a1950c · 2a1950c
1 parent 3997ee4
commit 2a1950c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/examples/language_model/transformer-xl/eval.py b/examples/language_model/transformer-xl/eval.py
@@ -40,7 +40,6 @@ def _evaluate(loader):
                 break
             ret = mem_transformer(src, target, *eval_mems)
             loss, eval_mems = ret[0], ret[1:]
-            seq_len = seq_len.numpy()
             eval_cur_loss = seq_len * loss.numpy()
             total_loss += eval_cur_loss
             total_len += seq_len

diff --git a/examples/language_model/transformer-xl/train.py b/examples/language_model/transformer-xl/train.py
@@ -242,9 +242,8 @@ def do_train(args):
                 logger.info(logger_info)
 
                 if args.save_model and rank == 0:
-                    model_dir = os.path.join(
-                        args.save_model,
-                        "step_" + str(step_idx) + "_" + str(eval_loss))
+                    model_dir = os.path.join(args.save_model,
+                                             "step_" + str(step_idx))
                     if not os.path.exists(model_dir):
                         os.makedirs(model_dir)
                     paddle.save(
@@ -253,6 +252,10 @@ def do_train(args):
                     paddle.save(
                         optimizer.state_dict(),
                         os.path.join(model_dir, "mem_transformer.pdopt"))
+                    f = open(
+                        os.path.join(args.save_model, "step_" + str(step_idx),
+                                     "evaluation_loss_" + str(eval_loss)), "w")
+                    f.close()
 
                 if args.scheduler == 'dev_perf':
                     scheduler.step(eval_loss)
@@ -271,6 +274,8 @@ def do_train(args):
 
                 mem_transformer.train()
 
+            if step_idx >= args.max_step:
+                return
             step_idx += 1
             batch_id += 1
             if args.scheduler in ['cosine', 'dev_perf']:
@@ -286,8 +291,6 @@ def do_train(args):
                     optimizer.set_lr(curr_lr)
             elif args.scheduler == 'noam':
                 scheduler.step()
-            if step_idx >= args.max_step:
-                return
 
     if args.save_model and rank == 0:
         model_dir = os.path.join(args.save_model, "step_final")