From c67a5d98d88abd729c928b05530131052e21bfc9 Mon Sep 17 00:00:00 2001
From: WeiXin <weixin10@baidu.com>
Date: Fri, 7 May 2021 09:17:25 +0800
Subject: [PATCH] pylayer_op:release context after compute. (#32707) (#32744)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

修复了py_layer_op由于没有析构PyLayerContext造成内存(显存)泄露的问题。

原始pr：#32707
---
 paddle/fluid/imperative/py_layer_fwd.h |  5 +++--
 paddle/fluid/operators/py_layer_op.cc  |  9 ++++++---
 paddle/fluid/operators/py_layer_op.h   | 11 +++++++++--
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/paddle/fluid/imperative/py_layer_fwd.h b/paddle/fluid/imperative/py_layer_fwd.h
index ccfd5b0e2db..de5f9d75e91 100644
--- a/paddle/fluid/imperative/py_layer_fwd.h
+++ b/paddle/fluid/imperative/py_layer_fwd.h
@@ -63,15 +63,16 @@ std::shared_ptr<GradOpNode> CreateGradOpNode(
   }
 }
 
-py::object PyLayerApply(const platform::Place& place, const py::object& cls,
+py::object PyLayerApply(const platform::Place& place, const py::handle& cls,
                         const py::args args, const py::kwargs kwargs) {
+  py::gil_scoped_acquire guard;
   auto bk_function = cls.attr("_backward_function");
   auto context = bk_function();
   auto forward = cls.attr("forward");
 
   auto result_forward = forward(context, *args, **kwargs);
   std::shared_ptr<operators::PyLayerContext> py_layer_ctx =
-      std::make_shared<operators::PyLayerContext>(context.release().ptr());
+      std::make_shared<operators::PyLayerContext>(context.ptr());
   // make inputs to varbase
   std::vector<std::shared_ptr<imperative::VarBase>> input_vars;
   // process args,`input_vars` only collect `imperative::VarBase`
diff --git a/paddle/fluid/operators/py_layer_op.cc b/paddle/fluid/operators/py_layer_op.cc
index 0090747d116..f91496eeab1 100644
--- a/paddle/fluid/operators/py_layer_op.cc
+++ b/paddle/fluid/operators/py_layer_op.cc
@@ -157,9 +157,12 @@ class PyLayerOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext &ctx) const override {
     auto &op_ = ctx.GetOp();
-    auto pylayer_op = dynamic_cast<const PyLayerOp *>(&op_);
-    if (pylayer_op) {
-      auto py_layer_context = pylayer_op->GetPyLayerContext();
+    auto const_pylayer_op = dynamic_cast<const PyLayerOp *>(&op_);
+    if (const_pylayer_op) {
+      auto pylayer_op = const_cast<PyLayerOp *>(const_pylayer_op);
+
+      // Release contex after executing the compute
+      auto py_layer_context = pylayer_op->ReleasePyLayerContext();
       py::object bk_ctx(py::handle(py_layer_context->GetMutableCtx()), true);
       auto &input_vars = ctx.MultiInputVar("X");
       auto output_vars = ctx.MultiOutputVar("Out");
diff --git a/paddle/fluid/operators/py_layer_op.h b/paddle/fluid/operators/py_layer_op.h
index 133435aa84d..d80faab90b2 100644
--- a/paddle/fluid/operators/py_layer_op.h
+++ b/paddle/fluid/operators/py_layer_op.h
@@ -34,6 +34,10 @@ class PyLayerContext {
   PyLayerContext() = delete;
 
   PyObject* GetMutableCtx() { return context_; }
+  ~PyLayerContext() {
+    py::gil_scoped_acquire guard;
+    Py_XDECREF(context_);
+  }
 
  private:
   PyObject* context_;
@@ -58,8 +62,11 @@ class PyLayerOp : public framework::OperatorWithKernel {
   void SetPyLayerContext(const std::shared_ptr<PyLayerContext>& py_context) {
     py_context_ = py_context;
   }
-  const std::shared_ptr<PyLayerContext>& GetPyLayerContext() const {
-    return py_context_;
+  std::shared_ptr<PyLayerContext> ReleasePyLayerContext() {
+    auto temp = py_context_;
+    py_context_.reset();
+    VLOG(3) << "`py_context_` in the PyLayerOp is released.";
+    return temp;
   }
 
  private:
-- 
GitLab