diff --git a/paddle/fluid/imperative/py_layer_fwd.h b/paddle/fluid/imperative/py_layer_fwd.h index bd132f2576fec1..873eab34bf23fa 100644 --- a/paddle/fluid/imperative/py_layer_fwd.h +++ b/paddle/fluid/imperative/py_layer_fwd.h @@ -63,15 +63,16 @@ std::shared_ptr 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 py_layer_ctx = - std::make_shared(context.release().ptr()); + std::make_shared(context.ptr()); // make inputs to varbase std::vector> 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 65e10181dcc3df..aad80de647a8be 100644 --- a/paddle/fluid/operators/py_layer_op.cc +++ b/paddle/fluid/operators/py_layer_op.cc @@ -151,9 +151,12 @@ class PyLayerOpKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext &ctx) const override { auto &op_ = ctx.GetOp(); - auto pylayer_op = dynamic_cast(&op_); - if (pylayer_op) { - auto py_layer_context = pylayer_op->GetPyLayerContext(); + auto const_pylayer_op = dynamic_cast(&op_); + if (const_pylayer_op) { + auto pylayer_op = const_cast(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 133435aa84d71e..d80faab90b2236 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& py_context) { py_context_ = py_context; } - const std::shared_ptr& GetPyLayerContext() const { - return py_context_; + std::shared_ptr ReleasePyLayerContext() { + auto temp = py_context_; + py_context_.reset(); + VLOG(3) << "`py_context_` in the PyLayerOp is released."; + return temp; } private: