Merge pull request #34 from chao9527/chao9527/PaddleBox

optimize for async
jack603047588 · Apr 1, 2022 · 3860903 · 3860903
2 parents 08d66d9 + 1588bcd
commit 3860903
Show file tree

Hide file tree

Showing 3 changed files with 252 additions and 124 deletions.
diff --git a/paddle/fluid/framework/boxps_trainer.cc b/paddle/fluid/framework/boxps_trainer.cc
@@ -139,18 +139,20 @@ void BoxPSTrainer::InitTrainerEnv(const ProgramDesc& main_program,
     }
   }
 
+  std::set<std::string> async_param_name;
   if (async_mode_) {
-    dense_table_->Init(*root_scope_, *param_need_sync_.get(),
+      async_param_name = dense_table_->Init(*root_scope_, *param_need_sync_.get(),
                        persistable_vars_);
   }
   for (int i = 0; i < thread_num_; ++i) {
     auto this_worker =
         std::dynamic_pointer_cast<paddle::framework::BoxPSWorker>(workers_[i]);
     this_worker->SetRootScope(root_scope_);
-    this_worker->CreateDeviceResource(main_program);
     if (async_mode_) {
-      this_worker->SetDenseTable(dense_table_.get());
+        this_worker->SetDenseTable(dense_table_.get());
+        this_worker->SetAsyncParamName(async_param_name);
     }
+    this_worker->CreateDeviceResource(main_program);
     //    CopyParameters(*root_scope_, i);
   }
 }