Customizable Python Layer in Dygraph

paddle/fluid/framework/operator.h

@@ -419,6 +419,7 @@ class ExecutionContext {
   const RuntimeContext Context() const { return ctx_; }
 
   std::string DebugString() const { return op_.DebugString(); }
+  const OperatorBase& GetOp() const { return op_; }
 
  private:
   const OperatorBase& op_;

paddle/fluid/imperative/dygraph_grad_maker.h

@@ -279,6 +279,8 @@ class TracedGradOp {
 
   void SetType(const std::string& type) { op_->SetType(type); }
 
+  const framework::OperatorBase& InnerOp() const { return op_->InnerOp(); }
+
   void SetAttrMap(const framework::AttributeMap& attrs) {
     return op_->SetAttrMap(attrs);
   }

paddle/fluid/imperative/layer.cc

@@ -406,7 +406,7 @@ void OpBase::Run(const framework::OperatorBase& op,
   OpBaseRunImpl<VariableWrapper>(op, ins, outs, attrs, place);
 }
 
-static void ClearNoNeedBufferInputs(OpBase* op) {
+void ClearNoNeedBufferInputs(OpBase* op) {
   auto& inferer = op->Info().NoNeedBufferVarsInferer();
   if (!inferer) return;
   auto* ins = op->GetMutableInsMap();

paddle/fluid/imperative/layer.h

@@ -286,5 +286,7 @@ std::shared_ptr<GradOpNode> CreateGradOpNode(
     const platform::Place& place,
     const std::map<std::string, std::string>& inplace_map);
 
+void ClearNoNeedBufferInputs(OpBase* op);
+
 }  // namespace imperative
 }  // namespace paddle

paddle/fluid/imperative/py_layer_fwd.h

@@ -0,0 +1,172 @@
+// 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 <string>
+#include <vector>
+#include "paddle/fluid/imperative/layer.h"
+#include "paddle/fluid/imperative/tracer.h"
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/type_defs.h"
+#include "paddle/fluid/operators/py_layer_op.h"
+
+namespace paddle {
+namespace imperative {
+
+namespace py = ::pybind11;
+
+bool RequiredGrad(const NameVarBaseMap& ins, const NameVarBaseMap& outs) {
+  for (const auto& name_pair : ins) {
+    for (const auto& var_base : name_pair.second) {
+      if (!var_base->OverridedStopGradient()) {
+        PassStopGradient(outs, var_base->OverridedStopGradient());
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+std::shared_ptr<GradOpNode> CreateGradOpNode(
+    const std::string& type, const NameVarBaseMap& ins,
+    const NameVarBaseMap& outs, const framework::AttributeMap& attrs,
+    const platform::Place& place,
+    const std::map<std::string, std::string>& inplace_map,
+    const std::shared_ptr<operators::PyLayerContext>& py_context) {
+  operators::PyLayerGradOpMaker<paddle::imperative::OpBase> maker(
+      type, ins, outs, attrs, inplace_map);
+
+  maker.SetPyLayerContext(py_context);
+  auto grad_node = maker();
+  if (grad_node && !grad_node->empty()) {
+    for (auto& grad_op : *grad_node) {
+      grad_op.SetId(OpBase::GenerateUniqueId());
+      grad_op.SetPlace(place);
+      ClearNoNeedBufferInputs(&grad_op);
+    }
+    return grad_node;
+  } else {
+    return nullptr;
+  }
+}
+
+py::object PyLayerApply(const platform::Place& place, const py::object& cls,
+                        const py::args args, const py::kwargs kwargs) {
+  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());
+  // make inputs to varbase
+  std::vector<std::shared_ptr<imperative::VarBase>> input_vars;
+  // process args,`input_vars` only collect `imperative::VarBase`
+  if (!args.empty()) {
+    for (auto ptr = args.begin(); ptr != args.end(); ptr++) {
+      try {
+        if (Py_None != ptr->ptr()) {
+          auto a = ptr->cast<std::shared_ptr<VarBase>>();
+          input_vars.push_back(a);
+        }
+      } catch (py::cast_error& err) {
+        // Only collect Tensor type in 'args' and pass them to backward. Ignore
+        // other types of input temporarily.
+      }
+    }
+  }
+  // process kwargs, only collect `imperative::VarBase`
+  if (!kwargs.empty()) {
+    for (auto ptr = kwargs.begin(); ptr != kwargs.end(); ptr++) {
+      try {
+        if (Py_None != ptr->second.ptr()) {
+          auto a = ptr->second.cast<std::shared_ptr<VarBase>>();
+          input_vars.push_back(a);
+        }
+      } catch (py::cast_error&) {
+        // Only collect Tensor type in 'kwargs' and pass them to backward.
+        // Ignore other types of input temporarily.
+      }
+    }
+  }
+  NameVarBaseMap ins = {{"X", input_vars}};
+
+  std::vector<std::shared_ptr<imperative::VarBase>> output_vars;
+  if (PyTuple_Check(result_forward.ptr()) ||
+      PyList_Check(result_forward.ptr())) {
+    auto tuple_result = result_forward.cast<py::tuple>();
+    for (size_t i = 0; i < tuple_result.size(); i++) {
+      if (Py_None != tuple_result[i].ptr()) {
+        try {
+          auto temp_out =
+              tuple_result[i].cast<std::shared_ptr<imperative::VarBase>>();
+          output_vars.push_back(temp_out);
+        } catch (py::cast_error&) {
+          PADDLE_THROW(platform::errors::Unimplemented(
+              "The output of `PyLayer.forward` should be `Tensor`."));
+        }
+      } else {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "The output of `PyLayer.forward` can not be `None`."));
+      }
+    }
+  } else {
+    if (Py_None != result_forward.ptr()) {
+      try {
+        auto temp_out =
+            result_forward.cast<std::shared_ptr<imperative::VarBase>>();
+        output_vars.push_back(temp_out);
+      } catch (py::cast_error&) {
+        PADDLE_THROW(platform::errors::Unimplemented(
+            "The output of `PyLayer.forward` should be `Tensor`."));
+      }
+    } else {
+      PADDLE_THROW(platform::errors::Unimplemented(
+          "The output of `PyLayer.forward` can not be `None`."));
+    }
+  }
+
+  NameVarBaseMap outs = {{"Out", output_vars}};
+
+  if (RequiredGrad(ins, outs)) {
+    std::map<std::string, std::string> inplace_map{};
+    bool if_inplace = false;
+    for (auto temp_ins : input_vars) {
+      if (if_inplace) {
+        break;
+      }
+      for (auto temp_outs : output_vars) {
+        if (temp_ins->Name() == temp_outs->Name()) {
+          if_inplace = true;
+          break;
+        }
+      }
+    }
+    if (if_inplace) {
+      inplace_map["X"] = "Out";
+    }
+
+    CreateGradOpNode("py_layer", ins, outs, {{}}, place, inplace_map,
+                     py_layer_ctx);
+  } else {
+    VLOG(3) << "No Grad to track for Op: py_layer_op";
+  }
+
+  return result_forward;
+}
+
+}  // namespace imperative
+}  // namespace paddle

paddle/fluid/imperative/tracer.cc

@@ -38,7 +38,7 @@ void SetCurrentTracer(const std::shared_ptr<Tracer>& tracer) {
   VLOG(6) << "Set current tracer: " << g_current_tracer;
 }
 
-static void PassStopGradient(const NameVarBaseMap& outs, bool generate_grad) {
+void PassStopGradient(const NameVarBaseMap& outs, bool generate_grad) {
   for (const auto& pair : outs) {
     for (const auto& var : pair.second) {
       // NOTE(zhiqiu): this happends when None output are passed from python

paddle/fluid/imperative/tracer.h

@@ -130,5 +130,7 @@ void IncreaseVarbaseReferenceCountUntilCopyComplete(
     const std::shared_ptr<imperative::VarBase>& var,
     const platform::Place& place);
 
+void PassStopGradient(const NameVarBaseMap& outs, bool generate_grad);
+
 }  // namespace imperative
 }  // namespace paddle

paddle/fluid/imperative/variable_wrapper.h

@@ -38,6 +38,9 @@ class VariableWrapper {
 
   explicit VariableWrapper(const std::string& name) : name_(name) {}
 
+  VariableWrapper(const std::string& name, const framework::Variable& variable)
+      : var_(variable), name_(name) {}
+
   ~VariableWrapper() { VLOG(10) << "Destruct VariableWrapper: " << Name(); }
 
   const framework::Variable& Var() const { return var_; }

paddle/fluid/operators/CMakeLists.txt

@@ -69,7 +69,7 @@ if(WITH_UNITY_BUILD)
     include(unity_build_rule.cmake)
 endif()
 
-register_operators(EXCLUDES py_func_op warpctc_op dgc_op lstm_op run_program_op eye_op recurrent_op
+register_operators(EXCLUDES py_layer_op py_func_op warpctc_op dgc_op lstm_op run_program_op eye_op recurrent_op
         sync_batch_norm_op ${OP_MKL_DEPS} DEPS ${OP_HEADER_DEPS})
 
 op_library(run_program_op SRCS run_program_op.cc run_program_op.cu.cc DEPS executor_cache ${OP_HEADER_DEPS})
@@ -161,6 +161,7 @@ endif()
 cc_library(tensor_formatter SRCS tensor_formatter.cc DEPS ${OP_HEADER_DEPS})
 if (WITH_PYTHON)
   cc_library(py_func_op SRCS py_func_op.cc DEPS op_registry python pybind)
+  cc_library(py_layer_op SRCS py_layer_op.cc DEPS op_registry python pybind)
 endif()
 
 set(GLOB_OP_LIB ${OP_LIBRARY} CACHE INTERNAL "Global OP library")