edit memory_sparse_table save thread_pool_size (PaddlePaddle#70)

danleifeng · Jul 22, 2022 · 933d14f · 933d14f
1 parent d05b5b5
commit 933d14f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/paddle/fluid/distributed/ps/table/memory_sparse_table.cc b/paddle/fluid/distributed/ps/table/memory_sparse_table.cc
@@ -35,14 +35,14 @@ namespace paddle {
 namespace distributed {
 
 int32_t MemorySparseTable::Initialize() {
-  _shards_task_pool.resize(_task_pool_size);
-  for (int i = 0; i < _shards_task_pool.size(); ++i) {
-    _shards_task_pool[i].reset(new ::ThreadPool(1));
-  }
   auto& profiler = CostProfiler::instance();
   profiler.register_profiler("pserver_sparse_update_all");
   profiler.register_profiler("pserver_sparse_select_all");
   InitializeValue();
+  _shards_task_pool.resize(_task_pool_size);
+  for (int i = 0; i < _shards_task_pool.size(); ++i) {
+    _shards_task_pool[i].reset(new ::ThreadPool(1));
+  }
   VLOG(0) << "initalize MemorySparseTable succ";
   return 0;
 }
@@ -58,9 +58,13 @@ int32_t MemorySparseTable::InitializeValue() {
     _real_local_shard_num =
         _real_local_shard_num < 0 ? 0 : _real_local_shard_num;
   }
+#ifdef PADDLE_WITH_HETERPS
+  _task_pool_size = _sparse_table_shard_num;
+#endif
   VLOG(1) << "memory sparse table _avg_local_shard_num: "
           << _avg_local_shard_num
-          << " _real_local_shard_num: " << _real_local_shard_num;
+          << " _real_local_shard_num: " << _real_local_shard_num
+          << " _task_pool_size:" << _task_pool_size;
 
   _local_shards.reset(new shard_type[_real_local_shard_num]);
 

diff --git a/paddle/fluid/distributed/ps/table/memory_sparse_table.h b/paddle/fluid/distributed/ps/table/memory_sparse_table.h
@@ -93,7 +93,7 @@ class MemorySparseTable : public Table {
   }
 
  protected:
-  const int _task_pool_size = 24;
+  int _task_pool_size = 24;
   int _avg_local_shard_num;
   int _real_local_shard_num;
   int _sparse_table_shard_num;