Improve handling of graph roots in autograd engine (pytorch#1635)

Author: apaszke

test/test_autograd.py

@@ -362,7 +362,7 @@ def hook(grad_input, grad_output):
         sum(fn(x, y)).sum().backward()
         self.assertTrue(was_called[0])
 
-    def _test_backward(self):
+    def test_backward(self):
         v_t = torch.randn(5, 5)
         x_t = torch.randn(5, 5)
         y_t = torch.rand(5, 5) + 0.1
@@ -385,9 +385,6 @@ def _test_backward(self):
         self.assertEqual(y.grad.data, y_grad * grad_output)
         self.assertEqual(z.grad.data, z_grad * grad_output)
 
-    def test_backward(self):
-        self._test_backward()
-
     def test_sparse_backward(self):
         class FixedGradientFunction(Function):
 
@@ -431,11 +428,6 @@ def backward(self, grad_x):
         (sparse_fn1(x) + sparse_fn2(x)).sum().backward()
         self.assertEqual(x.grad.data, sparse_grad1 + sparse_grad2)
 
-    @unittest.skip("BasicEngine is out of date")
-    def test_backward_basic_engine(self):
-        with backward_engine(torch.autograd.engine.BasicEngine):
-            self._test_backward()
-
     def test_multi_backward(self):
         x = Variable(torch.randn(5, 5), requires_grad=True)
         y = Variable(torch.randn(5, 5), requires_grad=True)
@@ -478,6 +470,18 @@ def test_multi_backward_no_grad(self):
         torch.autograd.backward([z, q], [torch.ones(5, 5), torch.ones(5, 5)])
         self.assertEqual(x.grad.data, torch.ones(5, 5))
 
+    def test_dependent_backward(self):
+        x = Variable(torch.randn(10), requires_grad=True)
+        y = x ** 2
+        z = y ** 3
+
+        go_y = torch.randn(10)
+        go_z = torch.randn(10)
+        torch.autograd.backward([y, z], [go_y, go_z])
+
+        xd = x.data
+        self.assertEqual(x.grad.data, 2 * xd * go_y + 6 * xd.pow(5) * go_z)
+
     def test_volatile(self):
         x = Variable(torch.ones(5, 5), requires_grad=True)
         y = Variable(torch.ones(5, 5) * 4, volatile=True)

torch/csrc/autograd/engine.cpp

@@ -1,4 +1,5 @@
 #include "torch/csrc/autograd/engine.h"
+#include "torch/csrc/autograd/functions/basic_ops.h"
 
 #include <atomic>
 #include <condition_variable>
@@ -226,9 +227,9 @@ auto Engine::evaluate_function(FunctionTask& task) -> void {
 }
 
 /** Finds all stochastic functions and appends them to the queue */
-auto Engine::find_stochastic_functions(function_queue& queue, GraphTask& task) -> void {
-  std::unordered_set<Function*> seen;
-  function_queue search_queue(queue);
+auto Engine::find_stochastic_functions(function_queue& queue, Function* graph_root, GraphTask& task) -> void {
+  std::unordered_set<Function*> seen {graph_root};
+  function_queue search_queue {graph_root};
   while (search_queue.size() > 0) {
     auto fn = search_queue.back(); search_queue.pop_back();
     for (auto& next_fn_pair : fn->next_functions) {
@@ -258,8 +259,6 @@ auto Engine::compute_dependencies(function_queue queue, GraphTask& task) -> void
   auto& dependencies = task.dependencies;
   while (queue.size() > 0) {
     auto fn = std::move(queue.back()); queue.pop_back();
-    // This is needed only to filter out roots that aren't executable
-    if (!fn->is_executable) continue;
     for (auto& next_fn_pair : fn->next_functions) {
       Function* next_ptr = next_fn_pair.first.get();
       if (!next_ptr) continue;
@@ -274,38 +273,6 @@ auto Engine::compute_dependencies(function_queue queue, GraphTask& task) -> void
   }
 }
 
-auto Engine::find_roots(const function_list& input_roots,
-                        variable_list& inputs,
-                        GraphTask& task) -> function_queue {
-  std::unordered_map<std::shared_ptr<Function>, std::unique_ptr<InputBuffer>> root_value;
-  int num_inputs = input_roots.size();
-  for (int i = 0; i < num_inputs; ++i) {
-    auto& input = inputs[i];
-    auto& root_info = input_roots[i];
-    auto root = root_info.first;
-    int input_nr = root_info.second;
-    auto& buf = root_value[root];
-    if (root->is_executable) {
-      if (!buf) buf.reset(new InputBuffer(root->num_inputs));
-      buf->add(input_nr, std::shared_ptr<Variable>(input));
-    }
-  }
-
-  function_queue roots;
-  for (auto& entry: root_value) {
-    const auto& root = entry.first;
-    roots.push_back(root.get());
-    // no need to enqueue tasks for non-executable functions
-    if (!root->is_executable) continue;
-    auto& input_buf = entry.second;
-    auto& queue = ready_queue(input_buf->device());
-    queue.push_front(FunctionTask(&task, root, std::move(*input_buf)));
-    task.has_any_work = true;
-  }
-
-  return roots;
-}
-
 auto Engine::execute(const function_list& input_roots,
                      variable_list& inputs,
                      bool keep_graph,
@@ -315,11 +282,19 @@ auto Engine::execute(const function_list& input_roots,
   GraphTask graph_task(keep_graph, callbacks);
   std::unique_lock<std::mutex> lock(graph_task.mutex);
 
-  // Find the unique roots and backprop into variables.
-  function_queue roots = find_roots(input_roots, inputs, graph_task);
+  auto graph_root = std::make_shared<GraphRoot>(input_roots, inputs);
+  function_queue roots;
+  for (auto entry : input_roots) {
+    if (entry.first->is_executable) {
+      graph_task.has_any_work = true;
+      roots.push_back(graph_root.get());
+      ready_queue(-1).push_front(FunctionTask(&graph_task, graph_root, InputBuffer(0)));
+      break;
+    }
+  }
 
   // Search the graph and find all stochastic functions. Append them to the queue.
-  find_stochastic_functions(roots, graph_task);
+  find_stochastic_functions(roots, graph_root.get(), graph_task);
 
   if (!graph_task.has_any_work) {
     throw std::runtime_error(

torch/csrc/autograd/engine.h

@@ -44,7 +44,7 @@ struct Engine {
       const function_list& roots,
       variable_list& inputs,
       GraphTask& task);
-  void find_stochastic_functions(function_queue& queue, GraphTask& task);
+  void find_stochastic_functions(function_queue& queue, Function* graph_root, GraphTask& task);
   void compute_dependencies(function_queue queue, GraphTask& task);
   void evaluate_function(FunctionTask& task);
   ReadyQueue& ready_queue(int device);

torch/csrc/autograd/functions/basic_ops.h

@@ -32,6 +32,20 @@ struct DelayedError : public Function {
   std::string msg;
 };
 
+struct GraphRoot : public Function {
+  GraphRoot(function_list functions, variable_list inputs)
+    : outputs(std::move(inputs)) {
+      next_functions = std::move(functions);
+      is_executable = true;
+    };
+
+  virtual variable_list apply(const variable_list& inputs) {
+    return outputs;
+  }
+
+  variable_list outputs;
+};
+
 struct Add : public Function {
   Add() {}