Merge pull request #1 from Yancey1989/panxin_polish_sparse

polish sparse update code
reyoung · Jun 8, 2018 · f25abba · f25abba
2 parents e0895e4 + 1239fce
commit f25abba
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 0 deletions.
diff --git a/paddle/fluid/operators/detail/request_handler_impl.cc b/paddle/fluid/operators/detail/request_handler_impl.cc
@@ -63,6 +63,9 @@ bool RequestSendHandler::Handle(const std::string& varname,
       PADDLE_THROW("sync: Can not find server side var");
       return false;
     }
+    if (invar->IsType<framework::SelectedRows>()) {
+      rpc_server_->RecordSparseVar(invar);
+    }
   }
 
   return true;

diff --git a/paddle/fluid/operators/detail/rpc_server.cc b/paddle/fluid/operators/detail/rpc_server.cc
@@ -73,6 +73,19 @@ void RPCServer::ResetBarrierCounter() {
     t.second = 0;
   }
 }
+void RPCServer::RecordSparseVar(framework::Variable* sparse_var) {
+  std::unique_lock<std::mutex> lock(mutex_sparse_var_recorder_);
+  sparse_vars_.push_back(sparse_var);
+}
+
+void RPCServer::ResetSparseVarsRecorder() {
+  VLOG(3) << "RPCServer reset sparse vars recorder.";
+  std::unique_lock<std::mutex> lock(mutex_sparse_var_recorder_);
+  for (auto* var : sparse_vars_) {
+    var->GetMutable<framework::SelectedRows>()->mutable_rows()->clear();
+  }
+  sparse_vars_.clear();
+}
 
 void RPCServer::RegisterRPC(const std::string& rpc_name,
                             RequestHandler* handler, int thread_num) {

diff --git a/paddle/fluid/operators/detail/rpc_server.h b/paddle/fluid/operators/detail/rpc_server.h
@@ -60,7 +60,10 @@ class RPCServer {
   void SetCond(const std::string& rpc_name);
   void WaitCond(const std::string& rpc_name);
   void IncreaseBatchBarrier(const std::string rpc_name);
+
   void ResetBarrierCounter();
+  void RecordSparseVar(framework::Variable* sparse_var);
+  void ResetSparseVarsRecorder();
 
  protected:
   virtual void ShutDownImpl() = 0;
@@ -74,6 +77,9 @@ class RPCServer {
   std::atomic<int> cur_cond_;
   std::condition_variable rpc_cond_;
 
+  std::vector<framework::Variable*> sparse_vars_;
+  std::mutex mutex_sparse_var_recorder_;
+
  protected:
   std::string bind_address_;
   std::atomic<int> exit_flag_;

diff --git a/paddle/fluid/operators/listen_and_serv_op.cc b/paddle/fluid/operators/listen_and_serv_op.cc
@@ -146,6 +146,7 @@ void ListenAndServOp::RunSyncLoop(framework::Executor *executor,
     rpc_service_->SetCond(detail::kRequestGet);
     rpc_service_->WaitBarrier(detail::kRequestGet);
     rpc_service_->ResetBarrierCounter();
+    rpc_service_->ResetSparseVarsRecorder();
   }  // while(true)
 }