upgrade lac (PaddlePaddle#969)

Author: joey12300

examples/lexical_analysis/train.py

@@ -41,6 +41,7 @@
 parser.add_argument("--max_seq_len", type=int, default=64, help="Number of words of the longest seqence.")
 parser.add_argument("--device", default="gpu", type=str, choices=["cpu", "gpu"] ,help="The device to select to train the model, is must be cpu/gpu.")
 parser.add_argument("--base_lr", type=float, default=0.001, help="The basic learning rate that affects the entire network.")
+parser.add_argument("--crf_lr", type=float, default=0.2, help="The learning rate ratio that affects CRF layers.")
 parser.add_argument("--emb_dim", type=int, default=128, help="The dimension in which a word is embedded.")
 parser.add_argument("--hidden_size", type=int, default=128, help="The number of hidden nodes in the GRU layer.")
 parser.add_argument("--logging_steps", type=int, default=10, help="Log every X updates steps.")
@@ -61,14 +62,19 @@ def evaluate(model, metric, data_loader):
         metric.update(num_infer_chunks.numpy(),
                       num_label_chunks.numpy(), num_correct_chunks.numpy())
         precision, recall, f1_score = metric.accumulate()
-    print("eval precision: %f, recall: %f, f1: %f" %
-          (precision, recall, f1_score))
+    logger.info("eval precision: %f, recall: %f, f1: %f" %
+                (precision, recall, f1_score))
     model.train()
+    return precision, recall, f1_score
 
 
 def train(args):
     paddle.set_device(args.device)
 
+    trainer_num = paddle.distributed.get_world_size()
+    if trainer_num > 1:
+        paddle.distributed.init_parallel_env()
+    rank = paddle.distributed.get_rank()
     # Create dataset.
     train_ds, test_ds = load_dataset(datafiles=(os.path.join(
         args.data_dir, 'train.tsv'), os.path.join(args.data_dir, 'test.tsv')))
@@ -117,24 +123,34 @@ def train(args):
         collate_fn=batchify_fn)
 
     # Define the model netword and its loss
-    model = BiGruCrf(args.emb_dim, args.hidden_size,
-                     len(word_vocab), len(label_vocab))
+    model = BiGruCrf(
+        args.emb_dim,
+        args.hidden_size,
+        len(word_vocab),
+        len(label_vocab),
+        crf_lr=args.crf_lr)
     # Prepare optimizer, loss and metric evaluator
     optimizer = paddle.optimizer.Adam(
         learning_rate=args.base_lr, parameters=model.parameters())
     chunk_evaluator = ChunkEvaluator(label_list=label_vocab.keys(), suffix=True)
 
     if args.init_checkpoint:
-        model_dict = paddle.load(args.init_checkpoint)
-        model.load_dict(model_dict)
-
+        if os.path.exists(args.init_checkpoint):
+            logger.info("Init checkpoint from %s" % args.init_checkpoint)
+            model_dict = paddle.load(args.init_checkpoint)
+            model.load_dict(model_dict)
+        else:
+            logger.info("Cannot init checkpoint from %s which doesn't exist" %
+                        args.init_checkpoint)
+    logger.info("Start training")
     # Start training
     global_step = 0
     last_step = args.epochs * len(train_loader)
     train_reader_cost = 0.0
     train_run_cost = 0.0
     total_samples = 0
     reader_start = time.time()
+    max_f1_score = -1
     for epoch in range(args.epochs):
         for step, batch in enumerate(train_loader):
             train_reader_cost += time.time() - reader_start
@@ -146,7 +162,7 @@ def train(args):
             train_run_cost += time.time() - train_start
             total_samples += args.batch_size
             if global_step % args.logging_steps == 0:
-                print(
+                logger.info(
                     "global step %d / %d, loss: %f, avg_reader_cost: %.5f sec, avg_batch_cost: %.5f sec, avg_samples: %.5f, ips: %.5f sequences/sec"
                     % (global_step, last_step, avg_loss, train_reader_cost /
                        args.logging_steps, (train_reader_cost + train_run_cost)
@@ -159,12 +175,21 @@ def train(args):
             optimizer.step()
             optimizer.clear_grad()
             if global_step % args.save_steps == 0 or global_step == last_step:
-                if paddle.distributed.get_rank() == 0:
-                    if args.do_eval:
-                        evaluate(model, chunk_evaluator, test_loader)
+                if rank == 0:
                     paddle.save(model.state_dict(),
                                 os.path.join(args.model_save_dir,
                                              "model_%d.pdparams" % global_step))
+                    logger.info("Save %d steps model." % (global_step))
+                    if args.do_eval:
+                        precision, recall, f1_score = evaluate(
+                            model, chunk_evaluator, test_loader)
+                        if f1_score > max_f1_score:
+                            max_f1_score = f1_score
+                            paddle.save(model.state_dict(),
+                                        os.path.join(args.model_save_dir,
+                                                     "best_model.pdparams"))
+                            logger.info("Save best model.")
+
             reader_start = time.time()