Merge pull request #15928 from velconia/imperative_backward_hooks

Imperative backward hooks

Author: velconia

paddle/fluid/framework/block_desc.cc

@@ -13,7 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/framework/block_desc.h"
+
 #include <queue>
+#include <unordered_set>
+#include <utility>
+
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/framework/program_desc.h"
 
@@ -155,6 +159,16 @@ void BlockDesc::RemoveOp(size_t s, size_t e) {
   ops_.erase(ops_.begin() + s, ops_.begin() + e);
 }
 
+void BlockDesc::RemoveOpInternal(const OpDesc *op_desc) {
+  // TODO(minqiyang): make this faster
+  for (auto it = ops_.begin(); it != ops_.end(); ++it) {
+    if (it->get() == op_desc) {
+      ops_.erase(it);
+      break;
+    }
+  }
+}
+
 std::vector<OpDesc *> BlockDesc::AllOps() const {
   std::vector<OpDesc *> res;
   for (const auto &op : ops_) {
@@ -163,20 +177,6 @@ std::vector<OpDesc *> BlockDesc::AllOps() const {
   return res;
 }
 
-void BlockDesc::Clear() {
-  // clear all ops
-  ops_.clear();
-
-  // clear all vars which are not persistable
-  for (auto it = vars_.begin(); it != vars_.end();) {
-    if (it->second->Persistable()) {
-      ++it;
-    } else {
-      vars_.erase(it++);
-    }
-  }
-}
-
 void BlockDesc::Flush() {
   for (auto &op_desc : ops_) {
     op_desc->Flush();

paddle/fluid/framework/block_desc.h

@@ -93,12 +93,12 @@ class BlockDesc {
    */
   void RemoveOp(size_t s, size_t e);
 
+  void RemoveOpInternal(const OpDesc *op_desc);
+
   void RemoveVar(const std::string &name) { vars_.erase(name); }
 
   std::vector<OpDesc *> AllOps() const;
 
-  void Clear();
-
   size_t OpSize() const { return ops_.size(); }
 
   OpDesc *Op(int idx) const { return ops_.at(idx).get(); }

paddle/fluid/framework/python_headers.h

@@ -24,3 +24,11 @@ limitations under the License. */
 
 #pragma pop_macro("_XOPEN_SOURCE")
 #pragma pop_macro("_POSIX_C_SOURCE")
+
+#if !defined(PYBIND11_HIDDEN)
+#ifdef _WIN32
+#define PYBIND11_HIDDEN __declspec(dllexport)
+#else
+#define PYBIND11_HIDDEN __attribute__((visibility("hidden")))
+#endif
+#endif

paddle/fluid/imperative/layer.cc

@@ -18,6 +18,7 @@
 #include <limits>
 #include <map>
 #include <random>
+#include <unordered_set>
 #include <utility>
 
 #include "paddle/fluid/framework/lod_tensor.h"
@@ -139,6 +140,8 @@ class Autograd {
           }
         }
       }
+
+      ready_op->InvokeBackwardHooks();
     }
   }
 
@@ -156,8 +159,10 @@ class Autograd {
       for (auto it : candidate->pre_ops_) {
         for (OpBase* pre_op : it.second) {
           if (!pre_op) continue;
-          VLOG(5) << "op dep " << candidate->op_desc_->Type() << " <---- "
-                  << it.first << " <---- " << pre_op->op_desc_->Type();
+          VLOG(5) << "op dep " << candidate->op_desc_->Type() << " trace id "
+                  << candidate->trace_id_ << " <---- " << it.first << " <---- "
+                  << pre_op->op_desc_->Type() << " trace id "
+                  << pre_op->trace_id_;
           if (visited.find(pre_op) == visited.end()) {
             visited.insert(pre_op);
             queue.push_back(pre_op);
@@ -211,6 +216,7 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
     return {};
   }
 
+  VLOG(3) << "apply op grad: " << op_desc_->Type();
   std::vector<framework::VariableValueMap> grad_outputs;
   if (backward_id_ > 0) {
     VLOG(3) << "py_layer_grad";
@@ -272,6 +278,22 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
   return input_vars_;
 }
 
+void OpBase::InvokeBackwardHooks() {
+  VLOG(3) << "call backward hooks, hooks num: " << backward_hooks_.size();
+
+  // call backward hooks
+  for (py::object& callable : backward_hooks_) {
+    callable(this);
+  }
+}
+
+void OpBase::RegisterBackwardHooks(const py::object& callable) {
+  VLOG(3) << "Register backward hooks " << trace_id_;
+
+  // TODO(minqiyang): check the callable format
+  backward_hooks_.push_back(callable);
+}
+
 void VarBase::RunBackward() {
   if (!pre_op_) return;
 

paddle/fluid/imperative/layer.h

@@ -123,22 +123,32 @@ class VarBase {
 
  private:
   VarBase(framework::Variable* var, VarBase* grad, bool stop_gradient)
-      : var_desc_(nullptr),
+      : name_(),
+        var_desc_(nullptr),
         var_(var),
         grads_(grad),
+        block_(nullptr),
+        persistable_(false),
         stop_gradient_(stop_gradient),
         pre_op_(nullptr),
+        pre_op_out_name_(),
         pre_op_out_idx_(-1) {}
 
  public:
   virtual ~VarBase() {
+    // TODO(minqiyang): remove var desc from block desc
     if (var_) {
       delete var_;
+      var_ = nullptr;
     }
 
     if (grads_) {
       delete grads_;
+      grads_ = nullptr;
     }
+
+    pre_op_ = nullptr;
+    pre_op_out_idx_ = -1;
   }
 
   inline OpBase* PreOp() const { return pre_op_; }
@@ -151,6 +161,14 @@ class VarBase {
 
   void RunBackward();
 
+  inline void ResetPreOp(OpBase* op) {
+    if (op == pre_op_) {
+      // clear pre_op info when op equals to var's pre_op
+      pre_op_ = nullptr;
+      pre_op_out_idx_ = -1;
+    }
+  }
+
   void TrackPreOp(OpBase* pre_op, const std::string& pre_op_out_name,
                   int pre_op_out_idx, bool pre_op_stop_gradient) {
     pre_op_ = pre_op;
@@ -184,11 +202,15 @@ class VarBase {
     return string::Sprintf("%s@IGrad", var_desc_->Name());
   }
 
+  std::string name_;
   framework::VarDesc* var_desc_;
 
   framework::Variable* var_;
   VarBase* grads_;
 
+  framework::BlockDesc* block_;
+  bool persistable_;
+
  private:
   bool stop_gradient_;
   OpBase* pre_op_;
@@ -199,22 +221,38 @@ class VarBase {
 /* The wrapper for OpDesc which holds a OpDesc and a OpDesc of its
  * gradient. This object should be managed totally by Python intepreter.
  */
-class OpBase {
+class PYBIND11_HIDDEN OpBase {
  public:
   OpBase()
       : op_desc_(nullptr),
         forward_id_(-1),
         backward_id_(-1),
-        place_(platform::CPUPlace()) {}
+        trace_id_(-1),
+        place_(platform::CPUPlace()),
+        backward_hooks_() {}
 
   virtual ~OpBase() {
+    // TODO(minqiyang): remove op_desc from block_desc in tracer
+    //
+    // reset all output vars' pre op
+    for (auto iter : output_vars_) {
+      for (VarBase* var : iter.second) {
+        var->ResetPreOp(this);
+      }
+    }
+
+    // release resource
     for (framework::OpDesc* desc : grad_op_descs_) {
       delete desc;
     }
   }
 
   std::map<std::string, std::vector<VarBase*>> ApplyGrad();
 
+  void RegisterBackwardHooks(const py::object& callable);
+
+  void InvokeBackwardHooks();
+
   // One of `op_desc_` or `forward_id_` is set, not both.
   // For pure python PyLayer, use `forward_id_`, otherwise, use op_desc_.
   framework::OpDesc* op_desc_;
@@ -225,6 +263,7 @@ class OpBase {
   // Note: each fwd op corresponds to a vector of bwd ops.
   std::vector<framework::OpDesc*> grad_op_descs_;
   int backward_id_;
+  int trace_id_;
 
   platform::Place place_;
 
@@ -239,6 +278,8 @@ class OpBase {
   std::vector<framework::VariableValueMap> grad_output_vars_;
 
   framework::BlockDesc* block_;
+
+  std::vector<py::object> backward_hooks_;
 };
 
 class Layer {

paddle/fluid/imperative/tracer.cc

@@ -14,7 +14,10 @@
 
 #include "paddle/fluid/imperative/tracer.h"
 
+#include <memory>
 #include <set>
+#include <unordered_map>
+#include <unordered_set>
 
 #include "paddle/fluid/operators/math/math_function.h"
 #include "paddle/fluid/platform/device_context.h"
@@ -110,7 +113,8 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
   std::map<std::string, VarBase*> vars;
 
   framework::OpDesc* op_desc = op->op_desc_;
-  VLOG(3) << "tracer tracing " << op_desc->Type();
+  VLOG(3) << "tracer tracing " << op_desc->Type() << " trace id "
+          << op->trace_id_;
   op_desc->InferShape(*block);
   op_desc->InferVarType(block);
 
@@ -133,11 +137,13 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
       if (inp->PreOp() && !inp->IsStopGradient()) {
         op->pre_ops_[it.first].push_back(inp->PreOp());
         op->pre_ops_out_idx_[it.first].push_back(inp->PreOpOutIdx());
+        VLOG(3) << "add pre op " << inp->PreOp()->op_desc_->Type();
       } else {
         op->pre_ops_[it.first].push_back(nullptr);
       }
       VLOG(3) << "input vname " << inp->var_desc_->Name() << " "
-              << inp->var_->IsInitialized();
+              << inp->var_->IsInitialized() << " stop_gradient "
+              << inp->IsStopGradient();
     }
   }
 
@@ -189,6 +195,7 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
 
     op->grad_input_vars_.resize(op->grad_op_descs_.size());
     op->grad_output_vars_.resize(op->grad_op_descs_.size());
+
     for (size_t i = 0; i < op->grad_op_descs_.size(); ++i) {
       framework::OpDesc* grad_op_desc = op->grad_op_descs_[i];
       for (auto it : grad_op_desc->Inputs()) {
@@ -201,7 +208,6 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
             PADDLE_ENFORCE(fwd_var_it != vars.end());
             // Forward inputs or outputs.
             grad_in_vars.push_back(fwd_var_it->second->var_);
-            vars_saved_for_backward.insert(it.first);
           } else {
             VarBase* var = vars[var_it->second];
             if (!var->grads_->var_->IsInitialized()) {
@@ -211,6 +217,8 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
             // Douts.
             grad_in_vars.push_back(var->grads_->var_);
           }
+
+          vars_saved_for_backward.insert(it.first);
         }
       }
 

paddle/fluid/pybind/imperative.h

@@ -33,7 +33,7 @@ class Layer : public imperative::Layer {
   }
 };
 
-class PyOpBase : public imperative::OpBase {
+class PYBIND11_HIDDEN PyOpBase : public imperative::OpBase {
  public:
   using imperative::OpBase::OpBase;  // Inherit constructors
 };

paddle/fluid/pybind/protobuf.cc

@@ -189,8 +189,6 @@ void BindBlockDesc(pybind11::module *m) {
              return self.HasVar(name);
            },
            pybind11::return_value_policy::reference)
-      .def("_clear_block", [](pd::BlockDesc &self) { return self.Clear(); },
-           pybind11::return_value_policy::reference)
       .def("_rename_var",
            [](pd::BlockDesc &self, const pybind11::bytes &byte_name,
               const pybind11::bytes &byte_name_new) {