Merge pull request #15292 from panyx0718/imperative

polish imperative codes

paddle/fluid/imperative/layer.cc

@@ -27,6 +27,9 @@
 namespace paddle {
 namespace imperative {
 
+const char* PyLayer::kFwdInp = "X";
+const char* PyLayer::kFwdOut = "Out";
+
 std::map<int, py::object> py_funcs_;
 
 using framework::Variable;
@@ -131,8 +134,9 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
   std::map<std::string, std::vector<framework::Variable*>> grad_outputs;
   if (backward_id_ > 0) {
     VLOG(3) << "py_layer_grad";
-    grad_outputs["Out@GRAD"] =
-        PyLayer::ApplyGrad(backward_id_, grad_input_vars_["X@GRAD"]);
+    grad_outputs[framework::GradVarName(PyLayer::kFwdOut)] = PyLayer::ApplyGrad(
+        backward_id_,
+        grad_input_vars_[framework::GradVarName(PyLayer::kFwdInp)]);
   } else {
     VLOG(3) << "op grad " << grad_op_desc_->Type();
     for (auto it : grad_output_vars_) {

paddle/fluid/imperative/layer.h

@@ -200,6 +200,9 @@ class PyLayer {
  public:
   virtual ~PyLayer() {}
 
+  static const char* kFwdInp;
+  static const char* kFwdOut;
+
   static void RegisterFunc(int func_id, const py::object& py_func);
 
   static int NumFuncs();

paddle/fluid/imperative/tracer.cc

@@ -164,28 +164,30 @@ std::vector<VarBase*> Tracer::PyTrace(OpBase* op,
                                       const std::vector<VarBase*>& inputs,
                                       bool stop_gradient) {
   VLOG(3) << "py_trace";
-  op->input_vars_["X"] = inputs;
-  op->output_vars_["Out"] = PyLayer::Apply(op->forward_id_, inputs);
+  op->input_vars_[PyLayer::kFwdInp] = inputs;
+  op->output_vars_[PyLayer::kFwdOut] = PyLayer::Apply(op->forward_id_, inputs);
   for (VarBase* inp : inputs) {
     if (inp->pre_op_) {
-      op->pre_ops_["X"].push_back(inp->pre_op_);
-      op->pre_ops_out_idx_["X"].push_back(inp->pre_op_out_idx_);
+      op->pre_ops_[PyLayer::kFwdInp].push_back(inp->pre_op_);
+      op->pre_ops_out_idx_[PyLayer::kFwdInp].push_back(inp->pre_op_out_idx_);
     } else {
-      op->pre_ops_["X"].push_back(nullptr);
+      op->pre_ops_[PyLayer::kFwdInp].push_back(nullptr);
     }
   }
 
-  auto& outputs = op->output_vars_["Out"];
+  auto& outputs = op->output_vars_[PyLayer::kFwdOut];
   for (size_t i = 0; i < outputs.size(); ++i) {
     VarBase* out = outputs[i];
     out->stop_gradient_ = stop_gradient;
     out->pre_op_ = op;
-    out->pre_op_out_name_ = "Out";
+    out->pre_op_out_name_ = PyLayer::kFwdOut;
     out->pre_op_out_idx_ = i;
   }
   if (!stop_gradient) {
-    auto& grad_input_vars = op->grad_input_vars_["X@GRAD"];
-    auto& grad_output_vars = op->grad_output_vars_["Out@GRAD"];
+    auto& grad_input_vars =
+        op->grad_input_vars_[framework::GradVarName(PyLayer::kFwdInp)];
+    auto& grad_output_vars =
+        op->grad_output_vars_[framework::GradVarName(PyLayer::kFwdOut)];
 
     for (const VarBase* inp : inputs) {
       grad_input_vars.push_back(inp->var_);

python/paddle/fluid/imperative/layers.py

@@ -54,6 +54,25 @@ class PyLayer(core.PyLayer):
     def __init__(self):
         super(PyLayer, self).__init__()
 
+    @classmethod
+    def _do_forward(cls, inputs):
+        return cls._to_tuple(cls.forward(inputs))
+
+    @classmethod
+    def _do_backward(cls, inputs):
+        return cls._to_tuple(cls.backward(inputs))
+
+    @staticmethod
+    def _to_tuple(inputs):
+        if not isinstance(inputs, list) and not isinstance(inputs, tuple):
+            inputs = [inputs]
+        ret = []
+        for inp in inputs:
+            tensor = core.LoDTensor()
+            tensor.set(inp, core.CPUPlace())
+            ret.append(tensor)
+        return tuple(ret)
+
     @staticmethod
     def forward(*inputs):
         raise NotImplementedError
@@ -70,16 +89,15 @@ class PyLayer(core.PyLayer):
 
         if not hasattr(cls, 'forward_id'):
             cls.forward_id = core.PyLayer.num_funcs() + 1
-            PyLayer.register_func(cls.forward_id, cls.forward)
+            PyLayer.register_func(cls.forward_id, cls._do_forward)
             cls.backward_id = core.PyLayer.num_funcs() + 1
-            PyLayer.register_func(cls.backward_id, cls.backward)
+            PyLayer.register_func(cls.backward_id, cls._do_backward)
 
         iop = core.OpBase()
         iop.forward_id = cls.forward_id
         iop.backward_id = cls.backward_id
         block.ops.append(iop)
         ivars = tracer.py_trace(iop, ivar_inputs, False)
-        # ivars = core.PyLayer.apply(cls.forward, inputs)
         ret = []
         for ivar in ivars:
             tensor = ivar.value().get_tensor()

python/paddle/fluid/tests/unittests/test_imperative.py

@@ -41,26 +41,12 @@ class MyPyLayer(fluid.imperative.PyLayer):
 
     @staticmethod
     def forward(inputs):
-        sys.stderr.write('before forward\n')
-        ret = np.tanh(inputs[0])
-        sys.stderr.write('after forward: %s\n' % ret)
-        tensor = core.LoDTensor()
-        tensor.set(ret, core.CPUPlace())
-        return tuple([tensor])
+        return np.tanh(inputs[0])
 
     @staticmethod
     def backward(inputs):
-        sys.stderr.write('calling into backward: %s\n' % str(inputs))
         inp, out, dout = inputs
-        inp = np.array(inp)
-        out = np.array(out)
-        dout = np.array(dout)
-        sys.stderr.write('calling into backward: %s, %s, %s\n' %
-                         (inp, out, dout))
-        ret = np.array(dout) * (1 - np.square(np.array(out)))
-        tensor = core.LoDTensor()
-        tensor.set(ret, core.CPUPlace())
-        return tuple([tensor])
+        return np.array(dout) * (1 - np.square(np.array(out)))
 
 
 class MLP(fluid.imperative.Layer):