Keep Reducer hooks in a vector instead of an unordered_map (#21783)

mrshenli · facebook-github-bot · commit da4e60226ccb · 2019-06-17T08:24:19.000-07:00
Summary: kuttas pointed out that the DDP Reducer only needs to remember `uintptr, Function` pairs, and hence does not need a nunordered map as added by #21591. Using a vector should speed it up a bit. Pull Request resolved: #21783 Differential Revision: D15854312 Pulled By: mrshenli fbshipit-source-id: 153ba035b8d658c7878a613f16a42de977d89c43
diff --git a/torch/csrc/distributed/c10d/reducer.cpp b/torch/csrc/distributed/c10d/reducer.cpp
@@ -105,14 +105,17 @@ Reducer::Reducer(
         auto grad_accumulator = variable.grad_accumulator();
 
         // Hook to execute after the gradient accumulator has executed.
-        hooks_[grad_accumulator->add_post_hook(
-            torch::make_unique<LambdaPostHook>([=] {
-                std::lock_guard<std::mutex> lock(this->mutex_);
-                this->mark_variable_ready(
-                    replica_index,
-                    variable_index,
-                    /* called_from_autograd= */ true);
-            }))] = grad_accumulator;
+        hooks_.emplace_back(
+            grad_accumulator->add_post_hook(
+                torch::make_unique<LambdaPostHook>([=] {
+                    std::lock_guard<std::mutex> lock(this->mutex_);
+                    this->mark_variable_ready(
+                        replica_index,
+                        variable_index,
+                        /* called_from_autograd= */ true);
+                })),
+            grad_accumulator
+        );
 
         // Map raw function pointer to replica index and parameter index.
         // This is used later on when the autograd graph is traversed
diff --git a/torch/csrc/distributed/c10d/reducer.h b/torch/csrc/distributed/c10d/reducer.h
@@ -54,7 +54,7 @@ class Reducer {
   std::vector<std::vector<std::shared_ptr<torch::autograd::Function>>>
       grad_accumulators_;
   std::unordered_map<torch::autograd::Function*, std::tuple<int, int>> func_;
-  std::unordered_map<uintptr_t, std::shared_ptr<torch::autograd::Function>>
+  std::vector<std::pair<uintptr_t, std::shared_ptr<torch::autograd::Function>>>
       hooks_;
 
   bool expect_autograd_hooks_;
