Support eager grad interface with allow_unused and multi startup_op

Author: veyron95

paddle/fluid/eager/accumulation/accumulation_node.cc

@@ -39,8 +39,8 @@ static void CopyOrAddTensor(paddle::experimental::Tensor* tensor,
 }
 
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeAccumulation::
-operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads) {
+operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+  const bool create_graph) {
  VLOG(3) << "Running Eager Backward Node: GradNodeAccumulation";
  PADDLE_ENFORCE(grads.size() == 1,
  paddle::platform::errors::Fatal(

paddle/fluid/eager/accumulation/accumulation_node.h

@@ -35,8 +35,8 @@ class GradNodeAccumulation : public GradNodeBase {
 
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads)
- override;
+ const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ const bool create_graph = false) override;
 
  std::string name() { return "GradNodeAccumulation"; }
 

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.cc

@@ -145,8 +145,8 @@ void GradNodeScale::SetTensorWrappers_X(
 void GradNodeScale::SetAttributes_scale(float scale) { scale_ = scale; }
 
 std::vector<std::vector<paddle::experimental::Tensor>> GradNodeScale::
-operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads) {
+operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+  const bool create_graph) {
  // 1. Check Output Size
  PADDLE_ENFORCE(
  ((grads.size() == 1) && (grads[0].size() == 1)),

paddle/fluid/eager/api/generated/eager_generated/backwards/scale_node.h

@@ -39,8 +39,8 @@ class GradNodeScale : public GradNodeBase {
 
  // Functor: perform backward computations
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads)
- override;
+ const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ const bool create_graph = false) override;
 
  void SetTensorWrappers_X(
  const std::vector<paddle::experimental::Tensor>& tensors);

paddle/fluid/eager/auto_code_generator/eager_generator.cc

@@ -2043,6 +2043,7 @@ static std::string GenerateGradNodeCCContents(
  }
 
  const char* BWD_RETURN_TEMPLATE =
+ " VLOG(1) << \"Run in GradNode, create_graph is: \" << create_graph; \n"
  " std::vector<std::vector<paddle::experimental::Tensor>> hooked_grads = "
  "GradNode%s::ApplyGradientHooks(grads);\n"
  " std::vector<std::vector<paddle::experimental::Tensor>> outputs(%d);\n"
@@ -2056,7 +2057,8 @@ static std::string GenerateGradNodeCCContents(
  const char* GRAD_FUNCTION_TEMPLATE =
  "std::vector<std::vector<paddle::experimental::Tensor>> "
  "GradNode%s::operator()(const "
- "std::vector<std::vector<paddle::experimental::Tensor>>& grads) {\n%s\n}";
+ "std::vector<std::vector<paddle::experimental::Tensor>>& grads, const "
+ "bool create_graph) {\n%s\n}";
  std::string grad_function_str = paddle::string::Sprintf(
  GRAD_FUNCTION_TEMPLATE, fwd_op_type, generated_grad_function_body);
 
@@ -2091,7 +2093,8 @@ static std::string GenerateGradNodeHeaderContents(
  "\n"
  " virtual std::vector<std::vector<paddle::experimental::Tensor>> "
  "operator()(const "
- "std::vector<std::vector<paddle::experimental::Tensor>>& grads) "
+ "std::vector<std::vector<paddle::experimental::Tensor>>& grads, const "
+ "bool create_graph = false) "
  "override;\n"
  "\n"
  " std::string name() override { return \" GradNode%s \"; } \n "

paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py

@@ -533,7 +533,7 @@ class {} : public egr::GradNodeBase {{
  ~{}() override = default;
 
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads) override;
+ const std::vector<std::vector<paddle::experimental::Tensor>>& grads, const bool create_graph = false) override;
  std::string name() override {{ return \" {} \"; }}
  // SetTensorWrapperX, SetTensorWrapperY, ...
  {}
@@ -609,7 +609,7 @@ def GenerateNodeDefinition(fwd_api_name, bwd_api_name, backward_fwd_input_map,
 
  grad_node_name = GetGradNodeName(fwd_api_name)
  FUNCTION_TEMPLATE = """
-std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads) {{
+std::vector<std::vector<paddle::experimental::Tensor>> {}::operator()(const std::vector<std::vector<paddle::experimental::Tensor>>& grads, const bool create_graph) {{
  // Call grad_api function
  auto grad_api_returns = paddle::experimental::{}({});
  {}

paddle/fluid/eager/backward.cc

@@ -34,6 +34,14 @@ std::unordered_map<GradNodeBase*, int> getInDegreeMap(
  // pass
  std::unordered_map<GradNodeBase*, int> node_in_degree_map;
 
+ // init potential startup node's indegree
+ std::queue<GradNodeBase*> queue_tmp = init_queue;
+ while (!queue_tmp.empty()) {
+ GradNodeBase* node = queue_tmp.front();
+ queue_tmp.pop();
+ node_in_degree_map[node] = 0;
+ }
+
  // Copy nodes
  std::queue<GradNodeBase*> queue = init_queue;
  std::unordered_set<GradNodeBase*> visited;
@@ -164,6 +172,7 @@ void GetGraphInfoBetweenTargets(
  }
  }
  }
+
  UpdateGraphInfo(target_nodes, depending_nodes, potential_stop_nodes);
 }
 
@@ -193,17 +202,33 @@ void GetTargetNodesInfo(const std::vector<paddle::experimental::Tensor>& inputs,
 
 std::vector<paddle::experimental::Tensor> GetResults(
  const std::vector<paddle::experimental::Tensor>& inputs,
- std::unordered_map<GradNodeBase*, paddle::experimental::Tensor>&
- result_map) {
+ const std::unordered_map<GradNodeBase*, paddle::experimental::Tensor>&
+ results_map,
+ bool allow_unused) {
  VLOG(1) << "Run in GetResults";
  if (inputs.empty()) return {};
 
  std::vector<paddle::experimental::Tensor> results;
  results.reserve(inputs.size());
- for (auto input : inputs) {
+ auto results_map_ = results_map;
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ auto& input = inputs[i];
  AutogradMeta* auto_grad_meta = EagerUtils::unsafe_autograd_meta(input);
  auto target_node = auto_grad_meta->GetMutableGradNode().get();
- results.emplace_back(result_map[target_node]);
+
+ if (results_map_.find(target_node) != results_map_.end()) {
+ // TODO(wuweilong): set StopGradient
+ // result_map[target_node].SetOverridedStopGradient(!create_graph_);
+ results.emplace_back(results_map_[target_node]);
+ } else {
+ PADDLE_ENFORCE_EQ(allow_unused, true,
+ paddle::platform::errors::InvalidArgument(
+ "The %d-th input does not appear in the backward "
+ "graph. Please check the input variable or set "
+ "allow_unused=True to get None result.",
+ i));
+ results.emplace_back();
+ }
  }
  return results;
 }
@@ -220,6 +245,20 @@ std::vector<paddle::experimental::Tensor> RunBackward(
  // *Inplace version check should perform at node-level
  // *Cross-batch accumulation happens at forward pass
 
+ /* --- Preprocess --- */
+
+ // TODO(wuweilong): output tensor duplicate check
+ // TODO(wuweilong): build no_grad_vars_grads according no_grad_vars
+ // TODO(wuweilong): output tensor' gradient is not in no_grad_vars
+
+ // TODO(wuweilong): check input tensor has grad op and stop_gradient = False
+ // TODO(wuweilong): input tensor duplicate check
+ // TODO(wuweilong): input tensor' gradient is not in no_grad_vars
+
+ // TODO(wuweilong): Prune output_targets which is not the input of startup_ops
+ // TODO(wuweilong): input == output case
+ // TODO(wuweilong): output_targets.size() should eaqul to output_grads.size()
+
  /* --- Initialization --- */
  // 1. Init queue with starting nodes
  // 2. Prepare initial input buffers
@@ -288,14 +327,28 @@ std::vector<paddle::experimental::Tensor> RunBackward(
  getInDegreeMap(queue);
 
  std::unordered_map<GradNodeBase*, AutogradMeta*> target_nodes_inputmeta_map;
- std::unordered_set<GradNodeBase*> target_nodes;
+ std::unordered_set<GradNodeBase*> target_nodes; // should be updated?
  GetTargetNodesInfo(inputs, &target_nodes, &target_nodes_inputmeta_map);
 
  std::unordered_map<GradNodeBase*, GradNodeBase*> depending_nodes;
  std::unordered_set<GradNodeBase*> potential_stop_nodes;
  GetGraphInfoBetweenTargets(queue, &target_nodes, &depending_nodes,
  &potential_stop_nodes);
 
+ std::unordered_set<GradNodeBase*> startup_ops_;
+ // ready_queue store all startup nodes
+ std::queue<GradNodeBase*> ready_queue;
+
+ // startup op's indegree should be 0
+ for (auto& pair : node_in_degree_map) {
+ if (pair.second == 0) {
+ auto* op = pair.first;
+ startup_ops_.emplace(op);
+ ready_queue.emplace(op);
+ }
+ }
+ VLOG(1) << " startup_ops' size is :" << startup_ops_.size();
+
  std::unordered_map<GradNodeBase*, paddle::experimental::Tensor> results_map;
 
  /* --- Topological Visit --- */
@@ -306,9 +359,9 @@ std::vector<paddle::experimental::Tensor> RunBackward(
  // |- Prepare for next node
  // 3. Update queue
  VLOG(6) << "Run Backward";
- while (!queue.empty()) {
- GradNodeBase* node = queue.front();
- queue.pop();
+ while (!ready_queue.empty()) {
+ GradNodeBase* node = ready_queue.front();
+ ready_queue.pop();
 
  // Run node: This is where Hook happens
  PADDLE_ENFORCE(
@@ -334,7 +387,7 @@ std::vector<paddle::experimental::Tensor> RunBackward(
 
  // Run Pre Backward Node and get outputs
  std::vector<std::vector<paddle::experimental::Tensor>> grad_output_tensors =
- (*node)(node_input_buffer->Buffers());
+ (*node)(node_input_buffer->Buffers(), create_graph);
  // TODO(jiabin): Should we erase it or find a more efficient way.
  node_input_buffers_dict.erase(node);
 
@@ -410,13 +463,13 @@ std::vector<paddle::experimental::Tensor> RunBackward(
  }
 
  if (node_in_degree_map[next_node] == 0 && !is_potential_stop_node) {
- queue.emplace(std::move(next_node));
+ ready_queue.emplace(std::move(next_node));
  }
  }
  }
  }
  if (!inputs.empty()) {
- return GetResults(inputs, results_map);
+ return GetResults(inputs, results_map, allow_unused);
  }
 
  VLOG(1) << "Run backward in the end, return {}";

paddle/fluid/eager/grad_node_info.h

@@ -95,7 +95,8 @@ class GradNodeBase {
  * is better choice to fit this format.
  * **/
  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads) = 0;
+ const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ const bool create_graph = false) = 0;
 
  /**
  * AddEdges is designed to set input tensors' backward Node as current

paddle/fluid/eager/tests/data_structure_tests/grad_node_test.h

@@ -32,8 +32,8 @@ class GradTestNode : public egr::GradNodeBase {
  GradTestNode() : GradNodeBase() { val_ = 1.0; }
  std::string name() override { return "GradTestNode"; }
  std::vector<std::vector<paddle::experimental::Tensor>> operator()(
- const std::vector<std::vector<paddle::experimental::Tensor>>& grads)
- override {
+ const std::vector<std::vector<paddle::experimental::Tensor>>& grads,
+ const bool create_graph = false) override {
  val_ = std::dynamic_pointer_cast<phi::DenseTensor>(grads[0][0].impl())
  ->data<float>()[0];
  phi::DenseTensorMeta meta =

python/paddle/fluid/tests/unittests/test_imperative_double_grad.py

@@ -19,6 +19,7 @@
 import unittest
 from unittest import TestCase
 import numpy as np
+import paddle.compat as cpt
 from paddle.fluid.framework import _test_eager_guard
 
 
@@ -61,6 +62,36 @@ def test_simple_example_eager_grad(self):
  with _test_eager_guard():
  self.func_simple_example_eager_grad()
  self.func_simple_example_eager_grad()
+ 
+ def func_simple_example_eager_grad_allow_unused(self):
+ np.random.seed(2021)
+ paddle.set_device('cpu')
+ np_x = np.random.random((3, 3))
+ np_y = np.random.random((3, 1))
+ np_z = np.random.random((3, 1))
+ x = paddle.to_tensor(np_x, dtype="float64", stop_gradient=False)
+ y = paddle.to_tensor(np_y, dtype="float64", stop_gradient=False)
+ z = paddle.to_tensor(np_z, dtype="float64", stop_gradient=False)
+ out_z = paddle.nn.functional.sigmoid(z)
+ out = paddle.matmul(x, y)
+
+ dx = fluid.dygraph.grad(out, [x, z], allow_unused=True)
+ 
+ dout = np.ones_like(np_y)
+ expected_dx = np.matmul(dout, np.transpose(np_y))
+ 
+ self.assertTrue(np.allclose(dx[0].numpy(), expected_dx[0]))
+ 
+ try:
+ dx = fluid.dygraph.grad(out, [x, z])
+ except ValueError as e:
+ error_msg = cpt.get_exception_message(e)
+ assert error_msg.find("allow_unused") > 0
+
+ def test_simple_example_eager_grad_allow_unused(self):
+ with _test_eager_guard():
+ self.func_simple_example_eager_grad_allow_unused()
+ self.func_simple_example_eager_grad_allow_unused()
 
 
 class TestDygraphDoubleGrad(TestCase):