Fix/Fix memory leak in dygraph (#17394)

JiabinYang · web-flow · commit d7df4e5e5b36 · 2019-05-17T14:09:11.000+08:00
* test=develop, add gradient sort backward strategy

* test=develop, fix test by add FLAGS_cudnn_deterministic on new tests

* test=develop, fix memory leak in dygraph mode

* test=develop, fix memory leak in dygraph mode

* test=develop, polish code

* test=develop, polish code

* test=develop, polish code
diff --git a/paddle/fluid/imperative/layer.cc b/paddle/fluid/imperative/layer.cc
@@ -112,8 +112,8 @@ void AddGradBySort(BackwardSumMap* bck_map, VarBase* target) {
         return a.first > b.first;
       });
   for (auto& var_pair : current.second) {
-    Variable* origin_grad = target->var_;
-    Variable* grad_to_add = var_pair.second->var_;
+    Variable* origin_grad = target->var_.get();
+    Variable* grad_to_add = var_pair.second->var_.get();
     VLOG(2) << "add origin_grad: " << target->Name();
     VLOG(2) << "added grad: " << var_pair.second->Name()
             << " trace id is: " << var_pair.first;
@@ -132,19 +132,19 @@ class Autograd {
       return;
     }
     VLOG(3) << "start autograd";
-    bck_map = new BackwardSumMap();
-    grad_ref = new GradientRef();
+    BackwardSumMap bck_map;
+    GradientRef grad_ref;
     std::deque<OpBase*> ready;
     ready.push_back(var->PreOp());
 
     std::map<OpBase*, int> dep_counts =
-        ComputeDepCounts(var->PreOp(), bck_stratedy);
+        ComputeDepCounts(var->PreOp(), bck_stratedy, &grad_ref);
 
     while (!ready.empty()) {
       OpBase* ready_op = ready.front();
       ready.pop_front();
       std::map<std::string, std::vector<VarBase*>> input_grads =
-          ready_op->ApplyGrad(bck_map, grad_ref, bck_stratedy);
+          ready_op->ApplyGrad(&bck_map, &grad_ref, bck_stratedy);
 
       for (auto it = input_grads.rbegin(); it != input_grads.rend(); ++it) {
         const std::vector<VarBase*>& ingrads = it->second;
@@ -171,7 +171,13 @@ class Autograd {
 
  private:
   std::map<OpBase*, int> ComputeDepCounts(
-      OpBase* op, const detail::BackwardStrategy& bck_stratedy) {
+      OpBase* op, const detail::BackwardStrategy& bck_stratedy,
+      GradientRef* grad_ref) {
+    if (bck_stratedy.sorted_sum_gradient_) {
+      PADDLE_ENFORCE_NOT_NULL(grad_ref,
+                              "grad_ref should not be null when "
+                              "using sorted grad backward strategy");
+    }
     std::map<OpBase*, int> ret;
 
     std::deque<OpBase*> queue;
@@ -185,13 +191,7 @@ class Autograd {
         for (const auto& map : candidate->grad_output_vars_) {
           for (const auto& it : map) {
             for (const auto& vb : it.second) {
-              if (grad_ref->find(vb) == grad_ref->end()) {
-                grad_ref->insert(std::make_pair(vb, 1));
-              } else {
-                // add ref count by 1 when we find grad_var can be generated by
-                // one grad_op
-                grad_ref->at(vb) += 1;
-              }
+              ++(*grad_ref)[vb];
             }
           }
         }
@@ -212,9 +212,6 @@ class Autograd {
     }
     return ret;
   }
-
-  BackwardSumMap* bck_map;
-  GradientRef* grad_ref;
 };
 
 std::unique_ptr<VarBase> VarBase::NewVarBase(const platform::Place& dst_place,
@@ -324,7 +321,7 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad(
           PADDLE_ENFORCE_NOT_NULL(grad_inp->var_, "op %s input %s nullptr",
                                   grad_op_desc->Type(), grad_inp->Name());
 
-          grad_invars.emplace_back(grad_inp->var_);
+          grad_invars.emplace_back(grad_inp->var_.get());
         }
       }
 
@@ -335,7 +332,7 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad(
           PADDLE_ENFORCE_NOT_NULL(grad_out->var_, "op %s output %s nullptr",
                                   grad_op_desc->Type(), grad_out->Name());
 
-          grad_outvars.emplace_back(grad_out->var_);
+          grad_outvars.emplace_back(grad_out->var_.get());
         }
       }
 
@@ -394,8 +391,8 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad(
             grad_ref->at(origin_outputs[i])--;
           }
         } else {
-          framework::Variable* grad = outputs[i]->var_;
-          framework::Variable* orig_grad = origin_outputs[i]->var_;
+          framework::Variable* grad = outputs[i]->var_.get();
+          framework::Variable* orig_grad = origin_outputs[i]->var_.get();
           VLOG(2) << "AddTo Called with orig_grad is: "
                   << origin_outputs[i]->name_ << " Grad to be added is "
                   << outputs[i]->name_;
@@ -451,7 +448,7 @@ void PyLayer::RegisterFunc(int func_id, const py::object& py_func) {
 
 int PyLayer::NumFuncs() { return py_funcs_.size(); }
 
-std::vector<framework::Variable*> PyLayer::Apply(
+std::vector<std::unique_ptr<framework::Variable>> PyLayer::Apply(
     int func_id, const std::vector<VarBase*>& inputs) {
   PADDLE_ENFORCE(py_funcs_.find(func_id) != py_funcs_.end());
   return CallPythonFunc(py_funcs_[func_id], inputs);
@@ -468,13 +465,13 @@ std::vector<VarBase*> PyLayer::ApplyGrad(int func_id,
     outs.emplace_back(new VarBase(
         string::Sprintf("%s_out_%d", framework::GradVarName(PyLayer::kFwdOut),
                         i),
-        rets[i], nullptr, true));
+        std::move(rets[i]), nullptr, true));
   }
 
   return outs;
 }
 
-std::vector<framework::Variable*> PyLayer::CallPythonFunc(
+std::vector<std::unique_ptr<framework::Variable>> PyLayer::CallPythonFunc(
     const py::object& callable, const std::vector<VarBase*>& ins) {
   py::gil_scoped_acquire guard;
   py::tuple in_args(ins.size());
@@ -488,19 +485,20 @@ std::vector<framework::Variable*> PyLayer::CallPythonFunc(
   auto ret = callable(in_args);
   auto ret_tuple = py::cast<py::tuple>(ret);
   size_t ret_num = py::len(ret_tuple);
-  std::vector<framework::Variable*> outs;
+  std::vector<std::unique_ptr<framework::Variable>> outs;
   outs.reserve(ret_num);
   VLOG(3) << "pyfunc out " << ret_num;
   for (size_t i = 0; i < ret_num; ++i) {
     try {
       auto* py_out_tensor = py::cast<framework::LoDTensor*>(ret_tuple[i]);
       PADDLE_ENFORCE_NOT_NULL(py_out_tensor,
                               "Output tensor %d should not be nullptr", i);
-      auto* var = new framework::Variable();
+      auto var =
+          std::unique_ptr<framework::Variable>(new framework::Variable());
       auto* tensor = var->GetMutable<framework::LoDTensor>();
       tensor->ShareDataWith(*py_out_tensor);
       tensor->set_lod(py_out_tensor->lod());
-      outs.emplace_back(var);
+      outs.emplace_back(std::move(var));
     } catch (py::cast_error&) {
       PADDLE_THROW("The %d-th output must be LoDTensor", i);
     }
diff --git a/paddle/fluid/imperative/layer.h b/paddle/fluid/imperative/layer.h
@@ -14,15 +14,16 @@
 
 #pragma once
 
-// clang-format off
-#include "paddle/fluid/framework/python_headers.h"
-// clang-format on
-
 #include <map>            // NOLINT
-#include <string>         // NOLINT
-#include <vector>         // NOLINT
 #include <memory>         // NOLINT
+#include <string>         // NOLINT
 #include <unordered_map>  // NOLINT
+#include <utility>
+#include <vector>  // NOLINT
+
+// clang-format off
+#include "paddle/fluid/framework/python_headers.h"
+// clang-format on
 
 #include "paddle/fluid/framework/op_desc.h"
 #include "paddle/fluid/framework/operator.h"
@@ -115,12 +116,14 @@ class OpBase;
 class VarBase {
  public:
   // Internal interface, create VarBase from exist variable
-  VarBase(const std::string& name, framework::Variable* var, VarBase* grad,
-          bool stop_gradient)
+  VarBase(const std::string& name, std::unique_ptr<framework::Variable> var,
+          VarBase* grad, bool stop_gradient)
       : VarBase(name, var->Get<framework::LoDTensor>().type(),
                 var->Get<framework::LoDTensor>().dims(),
-                var->Get<framework::LoDTensor>().place(), var, grad,
-                stop_gradient, false) {}
+                var->Get<framework::LoDTensor>().place(), nullptr, grad,
+                stop_gradient, false) {
+    var_ = std::move(var);
+  }
 
   // Python interface
   VarBase(const std::string& name, const framework::proto::VarType::Type dtype,
@@ -140,19 +143,19 @@ class VarBase {
   // TODO(minqiyang): need support SelectedRows
   VarBase(const std::string& name, framework::proto::VarType::Type dtype,
           const framework::DDim& shape, const platform::Place& place,
-          framework::Variable* var, VarBase* grad, bool stop_gradient,
-          bool persistable)
+          std::unique_ptr<framework::Variable> var, VarBase* grad,
+          bool stop_gradient, bool persistable)
       : name_(name),
         type_(framework::proto::VarType::LOD_TENSOR),
-        var_(var),
+        var_(std::move(var)),
         grads_(grad),
         stop_gradient_(stop_gradient),
         persistable_(persistable),
         pre_op_(nullptr),
         pre_op_out_name_(),
         pre_op_out_idx_(-1) {
     if (!var_) {
-      var_ = new framework::Variable();
+      var_.reset(new framework::Variable());
     }
     auto tensor = var_->GetMutable<framework::LoDTensor>();
     tensor->Resize(shape);
@@ -163,11 +166,6 @@ class VarBase {
 
  public:
   virtual ~VarBase() {
-    if (var_) {
-      delete var_;
-      var_ = nullptr;
-    }
-
     if (grads_) {
       delete grads_;
       grads_ = nullptr;
@@ -261,7 +259,7 @@ class VarBase {
   framework::proto::VarType::Type type_;
   platform::Place place_;
 
-  framework::Variable* var_;
+  std::unique_ptr<framework::Variable> var_;
   VarBase* grads_;
 
  private:
@@ -369,8 +367,8 @@ class Layer {
  public:
   virtual ~Layer() {}
 
-  virtual std::vector<VarBase> Forward(const std::vector<VarBase>& inputs) {
-    std::vector<VarBase> vars;
+  virtual std::vector<VarBase*> Forward(const std::vector<VarBase*>& inputs) {
+    std::vector<VarBase*> vars;
     return vars;
   }
 };
@@ -386,14 +384,14 @@ class PyLayer {
 
   static int NumFuncs();
 
-  static std::vector<framework::Variable*> Apply(
+  static std::vector<std::unique_ptr<framework::Variable>> Apply(
       int func_id, const std::vector<VarBase*>& inputs);
 
   static std::vector<VarBase*> ApplyGrad(int func_id,
                                          const std::vector<VarBase*>& inputs);
 
  private:
-  static std::vector<framework::Variable*> CallPythonFunc(
+  static std::vector<std::unique_ptr<framework::Variable>> CallPythonFunc(
       const py::object& callable, const std::vector<VarBase*>& ins);
 };
 
diff --git a/paddle/fluid/imperative/tracer.cc b/paddle/fluid/imperative/tracer.cc
@@ -18,6 +18,7 @@
 #include <set>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 
 #include "paddle/fluid/framework/var_type_inference.h"
 #include "paddle/fluid/operators/math/math_function.h"
@@ -153,7 +154,7 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
       PADDLE_ENFORCE_NOT_NULL(inp->var_, "op %s input %s nullptr", op->Type(),
                               inp->Name());
 
-      invars.emplace_back(inp->var_);
+      invars.emplace_back(inp->var_.get());
       if (!stop_gradient) {
         current_vars_map[inp->Name()] = inp;
       }
@@ -171,7 +172,7 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
     outvars.reserve(outputs.size());
     for (size_t i = 0U; i < outputs.size(); ++i) {
       VarBase* out = outputs[i];
-      outvars.emplace_back(out->var_);
+      outvars.emplace_back(out->var_.get());
       out->TrackPreOp(op, it.first, i, stop_gradient);
       if (!stop_gradient) {
         current_vars_map[out->Name()] = out;
@@ -294,17 +295,15 @@ std::vector<VarBase*> Tracer::PyTrace(OpBase* op,
 
   op->input_vars_[PyLayer::kFwdInp] = inputs;
 
-  std::vector<framework::Variable*> ret_vars =
+  std::vector<std::unique_ptr<framework::Variable>> ret_vars =
       PyLayer::Apply(op->forward_id_, inputs);
-
   op->TrackPreOp(PyLayer::kFwdInp, inputs);
 
   std::vector<VarBase*>& outputs = op->output_vars_[PyLayer::kFwdOut];
   outputs.reserve(ret_vars.size());
   for (size_t i = 0U; i != ret_vars.size(); ++i) {
-    framework::Variable* v = ret_vars[i];
-    VarBase* out = new VarBase(string::Sprintf("%s_out_%d", op->Type(), i), v,
-                               nullptr, stop_gradient);
+    VarBase* out = new VarBase(string::Sprintf("%s_out_%d", op->Type(), i),
+                               std::move(ret_vars[i]), nullptr, stop_gradient);
     outputs.emplace_back(out);
     out->TrackPreOp(op, PyLayer::kFwdOut, i, stop_gradient);
   }
diff --git a/paddle/fluid/pybind/imperative.h b/paddle/fluid/pybind/imperative.h
@@ -28,9 +28,9 @@ class Layer : public imperative::Layer {
  public:
   using imperative::Layer::Layer;  // Inherit constructors
 
-  std::vector<imperative::VarBase> Forward(
-      const std::vector<imperative::VarBase>& inputs) override {
-    PYBIND11_OVERLOAD(std::vector<imperative::VarBase>, Layer, Forward,
+  std::vector<imperative::VarBase*> Forward(
+      const std::vector<imperative::VarBase*>& inputs) override {
+    PYBIND11_OVERLOAD(std::vector<imperative::VarBase*>, Layer, Forward,
                       inputs);  // NOLINT
   }
 };
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -237,7 +237,8 @@ PYBIND11_MODULE(core, m) {
              return new_var.release();
            },
            py::return_value_policy::take_ownership)
-      .def("value", [](const imperative::VarBase &self) { return self.var_; },
+      .def("value",
+           [](const imperative::VarBase &self) { return self.var_.get(); },
            py::return_value_policy::reference)
       .def_property("name", &imperative::VarBase::Name,
                     &imperative::VarBase::SetName)
@@ -285,7 +286,7 @@ PYBIND11_MODULE(core, m) {
   py::class_<imperative::Layer, Layer /* <--- trampoline*/> layer(m, "Layer");
   layer.def(py::init<>())
       .def("forward", [](imperative::Layer &self,
-                         const std::vector<imperative::VarBase> &inputs) {
+                         const std::vector<imperative::VarBase *> &inputs) {
         return self.Forward(inputs);
       });
 
@@ -299,10 +300,9 @@ PYBIND11_MODULE(core, m) {
                 std::vector<imperative::VarBase *> outputs;
                 outputs.reserve(ret_vars.size());
                 for (size_t i = 0U; i != ret_vars.size(); ++i) {
-                  framework::Variable *v = ret_vars[i];
                   // TODO(minqiyang): use unique_name generator to set a name
-                  outputs.emplace_back(
-                      new imperative::VarBase("", v, nullptr, true));
+                  outputs.emplace_back(new imperative::VarBase(
+                      "", std::move(ret_vars[i]), nullptr, true));
                 }
 
                 return outputs;
