[Eager] Pylayer

paddle/fluid/eager/CMakeLists.txt

@@ -11,8 +11,9 @@ endif()
 add_subdirectory(api)
 add_subdirectory(accumulation)
 add_subdirectory(custom_operator)
-
-
+if(NOT ((NOT WITH_PYTHON) AND ON_INFER))
+    add_subdirectory(pylayer)
+endif()
 cc_library(grad_node_info SRCS grad_node_info.cc DEPS phi_api phi_tensor)
 cc_library(grad_tensor_holder SRCS grad_tensor_holder.cc DEPS grad_node_info gradient_accumulator)
 

paddle/fluid/eager/grad_node_info.cc

@@ -213,6 +213,49 @@ void GradNodeBase::SetGradInMeta(
   }
 }
 
+void GradNodeBase::SetGradInMeta(
+    const std::vector<paddle::experimental::Tensor*>& fwd_out,
+    size_t slot_rank) {
+  size_t slot_size = fwd_out.size();
+  PADDLE_ENFORCE_LE(
+      slot_rank, (bwd_in_meta_.size() - 1),
+      paddle::platform::errors::InvalidArgument(
+          "Slot Rank should less equal than bwd_in_meta_ size, since "
+          "bwd_in_meta_ is designed to hold as same num as backward "
+          "inputs."));
+  auto& metas = bwd_in_meta_.at(slot_rank);
+  // Init stop gradient vector before use to avoid push back
+  metas.resize(slot_size);
+  for (size_t i = 0; i < slot_size; i++) {
+    auto& meta = metas[i];
+    const auto& fwd_out_tensor = *fwd_out[i];
+    auto* fwd_out_meta =
+        egr::EagerUtils::nullable_autograd_meta(fwd_out_tensor);
+    PADDLE_ENFORCE_NOT_NULL(fwd_out_meta,
+                            paddle::platform::errors::PreconditionNotMet(
+                                "Bwd_in_meta should only be called while "
+                                "autograd_meta is not null. If you got this "
+                                "error, it indicates bugs in framework."));
+    if (fwd_out_meta->StopGradient()) {
+      // Set Stop Gradient only when its true or non-initialized autograd_meta,
+      // since all default value is false.
+      meta.SetStopGradient(fwd_out_meta->StopGradient());
+    }
+
+    // Record TensorMeta
+    if (phi::DenseTensor::classof(fwd_out_tensor.impl().get())) {
+      // Only Copy Meta
+      phi::DenseTensor* dense_tensor =
+          static_cast<phi::DenseTensor*>(fwd_out_tensor.impl().get());
+      meta.SetTensorMeta(dense_tensor->meta());
+      if (paddle::framework::IsComplexType(
+              paddle::framework::TransToProtoVarType(dense_tensor->type()))) {
+        need_complex_to_real_ = true;
+      }
+    }
+  }
+}
+
 void GradNodeBase::SetGradOutMeta(const paddle::experimental::Tensor& fwd_in,
                                   size_t slot_rank) {
   auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in);
@@ -300,6 +343,41 @@ void GradNodeBase::SetGradOutMeta(
   }
 }
 
+void GradNodeBase::SetGradOutMeta(
+    const std::vector<paddle::experimental::Tensor*>& fwd_in,
+    size_t slot_rank) {
+  size_t slot_size = fwd_in.size();
+  PADDLE_ENFORCE_LE(
+      slot_rank, (bwd_out_meta_.size() - 1),
+      paddle::platform::errors::InvalidArgument(
+          "Slot Rank should less equal than bwd_out_meta_ size, "
+          "since bwd_out_meta_ is designed to hold as same num as "
+          "backward outputs."));
+  auto& metas = bwd_out_meta_.at(slot_rank);
+  // Init stop gradient vector before use to avoid push back
+  metas.resize(slot_size);
+  for (size_t i = 0; i < slot_size; i++) {
+    const auto& fwd_in_tensor = *fwd_in[i];
+    auto& meta = metas[i];
+    auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in_tensor);
+    if (fwd_in_meta) {
+      // Set Stop Gradient only when its true or non-initialized autograd_meta,
+      // since all default value is false.
+      meta.SetStopGradient(fwd_in_meta->StopGradient());
+    }
+
+    // Record TensorMeta
+    if (fwd_in_tensor.impl() && fwd_in_tensor.impl().get()) {
+      if (phi::DenseTensor::classof(fwd_in_tensor.impl().get())) {
+        // Only Copy Meta
+        phi::DenseTensor* dense_tensor =
+            static_cast<phi::DenseTensor*>(fwd_in_tensor.impl().get());
+        meta.SetTensorMeta(dense_tensor->meta());
+      }
+    }
+  }
+}
+
 void GradNodeBase::SetDefaultGradInOutMeta() {
   PADDLE_ENFORCE((bwd_out_meta_.size() == 1) && (bwd_in_meta_.size() == 1),
                  paddle::platform::errors::PreconditionNotMet(

paddle/fluid/eager/grad_node_info.h

@@ -142,12 +142,15 @@ class GradNodeBase {
                      size_t slot_rank);
   void SetGradInMeta(const paddle::experimental::Tensor& fwd_out,
                      size_t slot_rank);
+  void SetGradInMeta(const std::vector<paddle::experimental::Tensor*>& fwd_out,
+                     size_t slot_rank);
 
   void SetGradOutMeta(const std::vector<paddle::experimental::Tensor>& fwd_in,
                       size_t slot_rank);
   void SetGradOutMeta(const paddle::experimental::Tensor& fwd_in,
                       size_t slot_rank);
-
+  void SetGradOutMeta(const std::vector<paddle::experimental::Tensor*>& fwd_in,
+                      size_t slot_rank);
   /**
    * Default setters for Grad in/out meta this should be used for same special
    * Node which will not create by user

paddle/fluid/eager/pylayer/CMakeLists.txt

@@ -0,0 +1 @@
+cc_library(py_layer_node SRCS py_layer_node.cc DEPS phi phi_api grad_node_info)

paddle/fluid/eager/pylayer/py_layer_node.cc

@@ -0,0 +1,159 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/eager/pylayer/py_layer_node.h"
+#include "paddle/fluid/eager/eager_tensor.h"
+
+#include "paddle/phi/api/all.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+#include "paddle/fluid/platform/device_context.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/fluid/platform/errors.h"
+#include "paddle/fluid/pybind/eager.h"
+#include "paddle/fluid/pybind/eager_utils.h"
+
+#include "glog/logging.h"
+#pragma GCC diagnostic ignored "-Wattributes"
+#include "pybind11/pytypes.h"
+
+namespace egr {
+std::vector<std::vector<paddle::experimental::Tensor>> GradNodePyLayer::
+operator()(
+    std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
+    bool create_graph) {
+  VLOG(3) << "Running Eager Backward Node: " << name();
+
+  std::vector<std::vector<paddle::experimental::Tensor>> hooked_grads =
+      GradNodePyLayer::ApplyGradientHooks(grads);
+
+  paddle::pybind::PyLayerObject* ctx =
+      reinterpret_cast<paddle::pybind::PyLayerObject*>(ctx_);
+
+  PADDLE_ENFORCE_EQ(ctx->forward_output_tensor_is_duplicable.size(),
+                    grads.size(),
+                    paddle::platform::errors::InvalidArgument(
+                        "%s's grad input size(%s) mast be equal with it's "
+                        "forward's output size(%s).",
+                        name(), grads.size(),
+                        ctx->forward_output_tensor_is_duplicable.size()));
+
+  auto backward_args = PyTuple_New(grads.size());
+  for (size_t i = 0; i < grads.size(); i++) {
+    if (ctx->forward_output_tensor_is_duplicable[i]) {
+      PyObject* pylist = PyList_New((Py_ssize_t)grads[i].size());
+      for (size_t j = 0; j < grads[i].size(); j++) {
+        if (ctx->materialize_grads && !grads[i][j].initialized()) {
+          paddle::experimental::Tensor tensor_tmp;
+          auto dense_tensor = std::make_shared<phi::DenseTensor>();
+          dense_tensor->set_meta(forward_outputs_meta_[i][j]);
+          tensor_tmp.set_impl(dense_tensor);
+          PyList_SET_ITEM(
+              pylist, static_cast<Py_ssize_t>(i),
+              paddle::pybind::ToPyObject(paddle::experimental::zeros_like(
+                  tensor_tmp, tensor_tmp.dtype(),
+                  forward_outputs_place_[i][j])));
+        } else {
+          PyList_SET_ITEM(pylist, static_cast<Py_ssize_t>(i),
+                          paddle::pybind::ToPyObject(grads[i][0], true));
+        }
+      }
+      PyTuple_SET_ITEM(backward_args, i, pylist);
+    } else {
+      if (ctx->materialize_grads && !grads[i][0].initialized()) {
+        paddle::experimental::Tensor tensor_tmp;
+        auto dense_tensor = std::make_shared<phi::DenseTensor>();
+        dense_tensor->set_meta(forward_outputs_meta_[i][0]);
+        tensor_tmp.set_impl(dense_tensor);
+        PyTuple_SET_ITEM(
+            backward_args, i,
+            paddle::pybind::ToPyObject(paddle::experimental::zeros_like(
+                tensor_tmp, tensor_tmp.dtype(), forward_outputs_place_[i][0])));
+      } else {
+        PyTuple_SET_ITEM(backward_args, i,
+                         paddle::pybind::ToPyObject(grads[i][0], true));
+      }
+    }
+  }
+
+  VLOG(6) << "PyLayer backward args is ready, begin call user's backward "
+             "function...";
+
+  auto backward_fn =
+      PyObject_GetAttrString(reinterpret_cast<PyObject*>(ctx), "backward");
+  if (!backward_fn) {
+    PADDLE_THROW(paddle::platform::errors::InvalidArgument(
+        "Get backward function faild."));
+  }
+  auto outputs = PyObject_CallObject(backward_fn, backward_args);
+  if (!outputs) {
+    PADDLE_THROW(paddle::platform::errors::External(
+        pybind11::detail::error_string().c_str()));
+  }
+
+  outputs_ = outputs;
+
+  VLOG(6) << "PyLayer backward function finish...";
+
+  PyObject* outputs_tuple = nullptr;
+  if (PyTuple_Check(outputs)) {
+    outputs_tuple = outputs;
+  } else {
+    outputs_tuple = PyTuple_New(1);
+    Py_INCREF(outputs);
+    PyTuple_SET_ITEM(outputs_tuple, 0, outputs);
+  }
+
+  size_t outputs_size = PyTuple_GET_SIZE(outputs_tuple);
+
+  if (outputs_size > ctx->forward_input_tensor_is_duplicable.size()) {
+    PADDLE_THROW(paddle::platform::errors::InvalidArgument(
+        "The number of outputs of `PyLayer.backward` should be %d, but "
+        "received %d.",
+        ctx->forward_input_tensor_is_duplicable.size(), outputs_size));
+  }
+
+  std::vector<std::vector<paddle::experimental::Tensor>> grad_out;
+  grad_out.reserve(ctx->forward_input_tensor_is_duplicable.size());
+  for (size_t i = 0; i < ctx->forward_input_tensor_is_duplicable.size(); i++) {
+    if (i < outputs_size) {
+      PyObject* obj = PyTuple_GET_ITEM(outputs_tuple, i);
+      if (this->OutputMeta()[i][0].IsStopGradient()) {
+        PADDLE_ENFORCE_EQ(
+            obj, Py_None,
+            paddle::platform::errors::InvalidArgument(
+                "%s's backward function should return None at %d position, "
+                "because it's forward Tensor's stopgradient is true.",
+                name(), i));
+        grad_out.push_back({});
+      } else {
+        if (ctx->forward_input_tensor_is_duplicable[i]) {
+          grad_out.push_back(paddle::pybind::GetTensorListFromPyObject(obj));
+        } else {
+          grad_out.push_back({paddle::pybind::GetTensorFromPyObject(obj)});
+        }
+      }
+    } else {
+      PADDLE_ENFORCE_EQ(
+          this->OutputMeta()[i][0].IsStopGradient(), true,
+          paddle::platform::errors::InvalidArgument(
+              "%s's backward function should not return empyt at %d position.",
+              name(), i));
+      grad_out.push_back({});
+    }
+  }
+
+  return grad_out;
+}
+}  // namespace egr

paddle/fluid/eager/pylayer/py_layer_node.h

@@ -0,0 +1,82 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Python.h>
+
+#include "paddle/fluid/eager/autograd_meta.h"
+#include "paddle/fluid/eager/grad_node_info.h"
+#include "paddle/fluid/eager/hooks.h"
+#include "paddle/phi/core/compat/convert_utils.h"
+#include "paddle/phi/core/tensor_meta.h"
+
+namespace egr {
+
+class GradNodePyLayer : public GradNodeBase {
+ public:
+  GradNodePyLayer(PyObject* ctx, size_t bwd_in_slot_num,
+                  size_t bwd_out_slot_num)
+      : GradNodeBase(bwd_in_slot_num, bwd_out_slot_num) {
+    ctx_ = ctx;
+  }
+
+  ~GradNodePyLayer() override { Py_DECREF(ctx_); };
+
+  virtual std::vector<std::vector<paddle::experimental::Tensor>> operator()(
+      std::vector<std::vector<paddle::experimental::Tensor>>& grads,  // NOLINT
+      bool create_graph = false) override;
+
+  void ClearTensorWrappers() override { VLOG(6) << "Do nothing here now"; }
+
+  bool IsTensorWrappersCleared() override {
+    VLOG(6) << "Do nothing here now";
+    return false;
+  }
+
+  std::string name() {
+    return "GradNodePyLayer_" + std::string(Py_TYPE(ctx_)->tp_name);
+  }
+
+  // for paddle.grad get result
+  PyObject* GetMutableOutputs() { return outputs_; }
+
+  void SaveForwardOutputsMeta(
+      const std::vector<std::vector<paddle::experimental::Tensor*>>&
+          outputs_tensor) {
+    forward_outputs_meta_.resize(outputs_tensor.size());
+    forward_outputs_place_.resize(outputs_tensor.size());
+    for (size_t i = 0; i < outputs_tensor.size(); i++) {
+      forward_outputs_meta_[i].reserve(outputs_tensor[i].size());
+      forward_outputs_place_[i].reserve(outputs_tensor[i].size());
+      for (auto tensor : outputs_tensor[i]) {
+        if (tensor->is_dense_tensor()) {
+          forward_outputs_meta_[i].push_back(
+              static_cast<phi::DenseTensor*>(tensor->impl().get())->meta());
+        } else {
+          forward_outputs_meta_[i].emplace_back();
+        }
+        forward_outputs_place_[i].emplace_back(tensor->inner_place());
+      }
+    }
+  }
+
+ private:
+  PyObject* ctx_{nullptr};
+  PyObject* outputs_{nullptr};
+  std::vector<std::vector<phi::DenseTensorMeta>> forward_outputs_meta_;
+  std::vector<std::vector<paddle::platform::Place>> forward_outputs_place_;
+};
+
+}  // namespace egr