Merge pull request #18 from qingshui/paddlebox

fix dataset read ins pipe
jack603047588 · Nov 9, 2021 · a8439c3 · a8439c3
2 parents 638d3f1 + 4e577e7
commit a8439c3
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 14 deletions.
diff --git a/paddle/fluid/framework/data_set.cc b/paddle/fluid/framework/data_set.cc
@@ -1515,6 +1515,14 @@ void PadBoxSlotDataset::CheckThreadPool(void) {
 void PadBoxSlotDataset::PreLoadIntoMemory() {
   CheckThreadPool();
   LoadIndexIntoMemory();
+  // dualbox global data shuffle
+  if (!FLAGS_padbox_dataset_disable_shuffle && mpi_size_ > 1) {
+    finished_counter_ = mpi_size_;
+    mpi_flags_.assign(mpi_size_, 1);
+    VLOG(3) << "RegisterClientToClientMsgHandler";
+    data_consumer_ = reinterpret_cast<void*>(new PadBoxSlotDataConsumer(this));
+    VLOG(3) << "RegisterClientToClientMsgHandler done";
+  }
 
   read_ins_ref_ = thread_num_;
   for (int64_t i = 0; i < thread_num_; ++i) {
@@ -1539,14 +1547,9 @@ void PadBoxSlotDataset::PreLoadIntoMemory() {
       }
     }));
   }
+
   // dualbox global data shuffle
   if (!FLAGS_padbox_dataset_disable_shuffle && mpi_size_ > 1) {
-    finished_counter_ = mpi_size_;
-    mpi_flags_.assign(mpi_size_, 1);
-    VLOG(3) << "RegisterClientToClientMsgHandler";
-    data_consumer_ = reinterpret_cast<void*>(new PadBoxSlotDataConsumer(this));
-    VLOG(3) << "RegisterClientToClientMsgHandler done";
-
     ShuffleData(shuffle_thread_num_);
     MergeInsKeys(shuffle_channel_);
   } else {
@@ -1577,6 +1580,14 @@ void PadBoxSlotDataset::LoadIntoMemory() {
 
   platform::Timer timeline;
   timeline.Start();
+  // dualbox global data shuffle
+  if (!FLAGS_padbox_dataset_disable_shuffle && mpi_size_ > 1) {
+    finished_counter_ = mpi_size_;
+    mpi_flags_.assign(mpi_size_, 1);
+    VLOG(3) << "RegisterClientToClientMsgHandler";
+    data_consumer_ = reinterpret_cast<void*>(new PadBoxSlotDataConsumer(this));
+    VLOG(3) << "RegisterClientToClientMsgHandler done";
+  }
 
   read_ins_ref_ = thread_num_;
   for (int64_t i = 0; i < thread_num_; ++i) {
@@ -1590,12 +1601,6 @@ void PadBoxSlotDataset::LoadIntoMemory() {
 
   // dualbox global data shuffle
   if (!FLAGS_padbox_dataset_disable_shuffle && mpi_size_ > 1) {
-    finished_counter_ = mpi_size_;
-    mpi_flags_.assign(mpi_size_, 1);
-    VLOG(3) << "RegisterClientToClientMsgHandler";
-    data_consumer_ = reinterpret_cast<void*>(new PadBoxSlotDataConsumer(this));
-    VLOG(3) << "RegisterClientToClientMsgHandler done";
-
     ShuffleData(shuffle_thread_num_);
     MergeInsKeys(shuffle_channel_);
   } else {

diff --git a/paddle/fluid/operators/pull_box_sparse_op.h b/paddle/fluid/operators/pull_box_sparse_op.h
@@ -22,6 +22,8 @@
 #include "paddle/fluid/framework/tensor.h"
 #include "paddle/fluid/operators/math/math_function.h"
 
+DECLARE_bool(enable_pull_box_padding_zero);
+
 namespace paddle {
 namespace operators {
 
@@ -123,8 +125,10 @@ static void PullBoxSparseFunctor(const framework::ExecutionContext &ctx) {
     const auto *slot = inputs[i];
     auto *output = outputs[i];
     if (slot->numel() == 0) {
-      // only support GPU
-      PaddingZeros<T>(ctx, output, batch_size, hidden_size);
+      if (FLAGS_enable_pull_box_padding_zero) {
+        // only support GPU
+        PaddingZeros<T>(ctx, output, batch_size, hidden_size);
+      }
       continue;
     }
     output->mutable_data<T>(ctx.GetPlace());

diff --git a/paddle/fluid/platform/flags.cc b/paddle/fluid/platform/flags.cc
@@ -604,3 +604,5 @@ DEFINE_bool(enable_pullpush_dedup_keys, false,
             "enable pull push dedup keys, default false");
 DEFINE_bool(enable_shuffle_by_searchid, false,
             "enable dualbox shuffle by searchid, default false");
+DEFINE_bool(enable_pull_box_padding_zero, true,
+            "enable pull box padding zero, default true");
diff --git a/python/paddle/fluid/__init__.py b/python/paddle/fluid/__init__.py
@@ -266,6 +266,7 @@ def __bootstrap__():
             'enable_slotrecord_reset_shrink',
             'enable_pullpush_dedup_keys',
             'enable_shuffle_by_searchid',
+            'enable_pull_box_padding_zero',
         ]
     core.init_gflags(["--tryfromenv=" + ",".join(read_env_flags)])
     core.init_glog(sys.argv[0])