add the time cost counter for the bert (#544)

Co-authored-by: Zeyu Chen <[email protected]>

Author: wawltor

examples/language_model/bert/run_pretrain.py

@@ -31,6 +31,7 @@
 from paddle.io import DataLoader, Dataset
 
 from paddlenlp.data import Stack, Tuple, Pad
+from paddlenlp.utils.tools import TimeCostAverage
 from paddlenlp.transformers import BertForPretraining, BertModel, BertPretrainingCriterion
 from paddlenlp.transformers import ErnieForPretraining, ErnieModel, ErniePretrainingCriterion
 from paddlenlp.transformers import BertTokenizer, ErnieTokenizer
@@ -377,13 +378,13 @@ def do_train(args):
             dataset_future = pool.submit(create_pretraining_dataset, data_file,
                                          args.max_predictions_per_seq,
                                          shared_file_list, args, worker_init)
-            train_reader_cost = 0.0
-            train_run_cost = 0.0
+            train_cost_avg = TimeCostAverage()
+            reader_cost_avg = TimeCostAverage()
             total_samples = 0
-            reader_start = time.time()
+            batch_start = time.time()
             for step, batch in enumerate(train_data_loader):
-                train_reader_cost += time.time() - reader_start
-                train_start = time.time()
+                train_reader_cost = time.time() - batch_start
+                reader_cost_avg.record(train_reader_cost)
                 global_step += 1
                 (input_ids, segment_ids, input_mask, masked_lm_positions,
                  masked_lm_labels, next_sentence_labels,
@@ -407,22 +408,23 @@ def do_train(args):
                     optimizer.step()
                 lr_scheduler.step()
                 optimizer.clear_grad()
-                train_run_cost += time.time() - train_start
                 total_samples += args.batch_size
+                train_run_cost = time.time() - batch_start
+                train_cost_avg.record(train_run_cost)
                 if global_step % args.logging_steps == 0:
                     if paddle.distributed.get_rank() == 0:
                         logger.info(
                             "global step: %d, epoch: %d, batch: %d, loss: %f, "
                             "avg_reader_cost: %.5f sec, avg_batch_cost: %.5f sec, avg_samples: %.5f, ips: %.5f sequences/sec"
                             % (global_step, epoch, step, loss,
-                               train_reader_cost / args.logging_steps,
-                               (train_reader_cost + train_run_cost) /
-                               args.logging_steps, total_samples /
-                               args.logging_steps, total_samples /
-                               (train_reader_cost + train_run_cost)))
-                    train_reader_cost = 0.0
-                    train_run_cost = 0.0
+                               reader_cost_avg.get_average(),
+                               train_cost_avg.get_average(), total_samples /
+                               args.logging_steps, total_samples / (
+                                   args.logging_steps *
+                                   train_cost_avg.get_average())))
                     total_samples = 0
+                    train_cost_avg.reset()
+                    reader_cost_avg.reset()
                 if global_step % args.save_steps == 0:
                     if paddle.distributed.get_rank() == 0:
                         output_dir = os.path.join(args.output_dir,
@@ -440,7 +442,7 @@ def do_train(args):
                 if global_step >= args.max_steps:
                     del train_data_loader
                     return
-                reader_start = time.time()
+                batch_start = time.time()
 
             del train_data_loader
             train_data_loader, data_file = dataset_future.result(timeout=None)

examples/language_model/bert/static/run_pretrain.py

@@ -29,6 +29,7 @@
 import paddle.distributed.fleet as fleet
 from paddle.io import DataLoader, Dataset
 
+from paddlenlp.utils.tools import TimeCostAverage
 from paddlenlp.transformers import BertForPretraining, BertModel, BertPretrainingCriterion
 from paddlenlp.transformers import BertTokenizer
 from paddlenlp.transformers import LinearDecayWithWarmup
@@ -367,34 +368,35 @@ def do_train(args):
                                          data_holders, worker_init,
                                          paddle.static.cuda_places())
 
-            train_reader_cost = 0.0
-            train_run_cost = 0.0
+            train_cost_avg = TimeCostAverage()
+            reader_cost_avg = TimeCostAverage()
             total_samples = 0
-            reader_start = time.time()
+            batch_start = time.time()
             for step, batch in enumerate(train_data_loader):
-                train_reader_cost += time.time() - reader_start
+                train_reader_cost = time.time() - batch_start
+                reader_cost_avg.record(train_reader_cost)
                 global_step += 1
                 train_start = time.time()
                 loss_return = exe.run(main_program,
                                       feed=batch,
                                       fetch_list=[loss])
-                train_run_cost += time.time() - train_start
                 total_samples += args.batch_size
                 # In the new 2.0 api, must call this function to change the learning_rate
                 lr_scheduler.step()
+                train_run_cost = time.time() - batch_start
+                train_cost_avg.record(train_run_cost)
                 if global_step % args.logging_steps == 0:
                     print(
                         "tobal step: %d, epoch: %d, batch: %d, loss: %f, "
                         "avg_reader_cost: %.5f sec, avg_batch_cost: %.5f sec, avg_samples: %.5f, ips: %.5f sequences/sec"
-                        %
-                        (global_step, epoch, step, loss_return[0],
-                         train_reader_cost / args.logging_steps,
-                         (train_reader_cost + train_run_cost) /
-                         args.logging_steps, total_samples / args.logging_steps,
-                         total_samples / (train_reader_cost + train_run_cost)))
-                    train_reader_cost = 0.0
-                    train_run_cost = 0.0
+                        % (global_step, epoch, step, loss_return[0],
+                           reader_cost_avg.get_average(),
+                           train_cost_avg.get_average(),
+                           total_samples / args.logging_steps, total_samples /
+                           (args.logging_steps * train_cost_avg.get_average())))
                     total_samples = 0
+                    train_cost_avg.reset()
+                    reader_cost_avg.reset()
                 if global_step % args.save_steps == 0:
                     if worker_index == 0:
                         output_dir = os.path.join(args.output_dir,
@@ -410,7 +412,7 @@ def do_train(args):
                     reader_start = time.time()
                     del train_data_loader
                     return
-                reader_start = time.time()
+                batch_start = time.time()
             del train_data_loader
             train_data_loader, data_file = dataset_future.result(timeout=None)
         epoch += 1

paddlenlp/utils/tools.py

@@ -57,3 +57,34 @@ def dygraph_params_to_static(model, dygraph_tensor_dict):
         ret_dict[parm.name] = dygraph_tensor_dict[name]
 
     return ret_dict
+
+
+class TimeCostAverage(object):
+    """
+    Simple tool for calcluating time average cost in the process of training and inferencing.
+    """
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        """
+        Reset the recoder state, and reset the `cnt` to zero.
+        """
+        self.cnt = 0
+        self.total_time = 0
+
+    def record(self, usetime):
+        """
+        Recoding the time cost in current step and accumulating the `cnt`.
+        """
+        self.cnt += 1
+        self.total_time += usetime
+
+    def get_average(self):
+        """
+        Returning the average time cost after the start of training.
+        """
+        if self.cnt == 0:
+            return 0
+        return self.total_time / self.cnt