support multi process for bert

BERT/reader/cls.py

@@ -18,7 +18,7 @@ import csv
 import numpy as np
 import tokenization
 from batching import prepare_batch_data
-import functools
+
 
 class DataProcessor(object):
     """Base class for data converters for sequence classification data sets."""
@@ -178,37 +178,16 @@ class DataProcessor(object):
                 yield batch, total_token_num
 
         def wrapper():
-            trainers_num = int(os.environ.get('PADDLE_TRAINERS_NUM', 1))
-            trainer_id = int(os.getenv("PADDLE_TRAINER_ID", 0)) + 1
-            if trainers_num > 1:
-                print("start data reader (trainers_num: {}, trainer_id: {})".format(
-                    trainers_num, trainer_id-1))
-            get_prepared_batch_input = functools.partial(
-               self.generate_batch_data,
+            for batch_data, total_token_num in batch_reader(
+                    instance_reader, batch_size, self.in_tokens):
+                batch_data = self.generate_batch_data(
+                    batch_data,
+                    total_token_num,
                     voc_size=-1,
                     mask_id=-1,
                     return_input_mask=True,
                     return_max_len=False,
                     return_num_token=False)
-
-            train_data, train_token_num, idx = None, None, 1
-            for batch_data, total_token_num in batch_reader(
-                    instance_reader, batch_size, self.in_tokens):
-                if trainers_num > 1:
-                    if idx < trainers_num:
-                        if idx == trainer_id:
-                            train_data, train_token_num = batch_data, total_token_num
-                        idx += 1
-                    else:
-                        if idx == trainer_id:
-                            train_data, train_token_num = batch_data, total_token_num
-                        assert train_data is not None, "train data should not be None."
-                        assert train_token_num is not None, "train data should not be None."
-                        batch_data = get_prepared_batch_input(train_data, train_token_num)
-                        yield batch_data
-                        train_data, train_token_num, idx = None, None, 1
-                else:
-                    batch_data = get_prepared_batch_input(batch_data, total_token_num)
                 yield batch_data
 
         return wrapper

BERT/run_classifier.py

@@ -278,6 +278,10 @@ def main(args):
             exec_strategy=exec_strategy,
             build_strategy = build_strategy,
             main_program=train_program)
+        num_trainers = int(os.environ.get('PADDLE_TRAINERS_NUM', 1))
+        if num_trainers > 1:
+            train_data_generator = fluid.contrib.reader.multi_process_reader(
+                train_data_generator)
                         
         train_pyreader.decorate_tensor_provider(train_data_generator)
     else: