From 11d4d39cd7ce5dd33df94ed65854eb357cad2400 Mon Sep 17 00:00:00 2001
From: Xin Pan <panxin.grad@gmail.com>
Date: Tue, 8 Jan 2019 13:16:45 +0800
Subject: [PATCH] forward working

test=develop
---
 paddle/fluid/imperative/layer.h               | 21 +++++++++++--------
 paddle/fluid/pybind/pybind.cc                 |  5 +++--
 python/paddle/fluid/framework.py              |  5 ++++-
 python/paddle/fluid/imperative/layers.py      | 16 ++++++++++++--
 .../fluid/tests/unittests/test_imperative.py  |  4 ++--
 5 files changed, 35 insertions(+), 16 deletions(-)

diff --git a/paddle/fluid/imperative/layer.h b/paddle/fluid/imperative/layer.h
index d76512af0..a0eee357c 100644
--- a/paddle/fluid/imperative/layer.h
+++ b/paddle/fluid/imperative/layer.h
@@ -161,13 +161,14 @@ class Layer {
 
 static void CallPythonFunc(py::object* callable,
                            const std::vector<framework::LoDTensor>& ins,
-                           std::vector<framework::LoDTensor*>* outs) {
+                           std::vector<VarBase*>* outs) {
   py::gil_scoped_acquire guard;
   py::tuple in_args(ins.size());
   for (size_t i = 0; i < ins.size(); ++i) {
     in_args[i] = ins[i].IsInitialized() ? py::cast(ins[i]) : py::cast(nullptr);
   }
 
+  // TODO(panyx0718): Who owns the returned LoDTensor.
   auto ret = (*callable)(in_args);
   auto ret_tuple = py::cast<py::tuple>(ret);
   size_t ret_num = py::len(ret_tuple);
@@ -176,7 +177,11 @@ static void CallPythonFunc(py::object* callable,
       auto* py_out_tensor = py::cast<framework::LoDTensor*>(ret_tuple[i]);
       PADDLE_ENFORCE_NOT_NULL(py_out_tensor,
                               "Output tensor %d should not be nullptr", i);
-      outs->push_back(py_out_tensor);
+      VarBase* var = new VarBase();
+      auto* tensor = var->var_->GetMutable<framework::LoDTensor>();
+      tensor->ShareDataWith(*py_out_tensor);
+      tensor->set_lod(py_out_tensor->lod());
+      outs->push_back(var);
     } catch (py::cast_error&) {
       PADDLE_THROW("The %d-th output must be LoDTensor", i);
     }
@@ -187,18 +192,16 @@ class PyLayer {
  public:
   virtual ~PyLayer() {}
 
-  static std::vector<VarBase> Apply(py::object* callable,
-                                    const std::vector<VarBase>& inputs) {
-    std::vector<VarBase> outputs;
+  static std::vector<VarBase*> Apply(py::object* callable,
+                                     const std::vector<VarBase>& inputs) {
     std::vector<framework::LoDTensor> tensor_inputs;
-    std::vector<framework::LoDTensor*> tensor_outputs;
+    std::vector<VarBase*> ret;
 
     for (const VarBase& in : inputs) {
       tensor_inputs.push_back(in.var_->Get<framework::LoDTensor>());
     }
-
-    CallPythonFunc(callable, tensor_inputs, &tensor_outputs);
-    return outputs;
+    CallPythonFunc(callable, tensor_inputs, &ret);
+    return ret;
   }
 };
 
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index d065818bc..684b931ee 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -182,9 +182,10 @@ PYBIND11_MODULE(core, m) {
       .def_static("apply",
                   [](py::object *callable,
                      const std::vector<imperative::VarBase> &inputs)
-                      -> std::vector<imperative::VarBase> {
+                      -> std::vector<imperative::VarBase *> {
                         return imperative::PyLayer::Apply(callable, inputs);
-                      });
+                      },
+                  py::return_value_policy::take_ownership);
 
   BindTracer(&m);
 
diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py
index 70767c962..8e18dffac 100644
--- a/python/paddle/fluid/framework.py
+++ b/python/paddle/fluid/framework.py
@@ -372,7 +372,10 @@ class Variable(object):
         self.stop_gradient = stop_gradient
         self.is_data = is_data
         if _in_imperative_mode():
-            self._ivar = core.VarBase()
+            if 'ivar' in kwargs:
+                self._ivar = kwargs['ivar']
+            else:
+                self._ivar = core.VarBase()
             self._ivar.desc = self.desc
             self._ivar.stop_gradient = stop_gradient
 
diff --git a/python/paddle/fluid/imperative/layers.py b/python/paddle/fluid/imperative/layers.py
index 769642736..06b6d7ac0 100644
--- a/python/paddle/fluid/imperative/layers.py
+++ b/python/paddle/fluid/imperative/layers.py
@@ -67,5 +67,17 @@ class PyLayer(core.PyLayer):
     def __call__(cls, inputs):
         inputs = map(base.to_variable, inputs)
         inputs = [x._ivar for x in inputs]
-        sys.stderr.write('%s\n' % inputs)
-        return core.PyLayer.apply(cls.forward, inputs)
+        ivars = core.PyLayer.apply(cls.forward, inputs)
+        ret = []
+        for ivar in ivars:
+            tensor = ivar.value.get_tensor()
+            block = framework.default_main_program().current_block()
+            py_var = framework.Variable(
+                block,
+                type=core.VarDesc.VarType.LOD_TENSOR,
+                name=None,
+                shape=tensor.shape(),
+                dtype=tensor._dtype(),
+                ivar=ivar)
+            ret.append(py_var)
+        return ret
diff --git a/python/paddle/fluid/tests/unittests/test_imperative.py b/python/paddle/fluid/tests/unittests/test_imperative.py
index ae99fb82e..133e1e65c 100644
--- a/python/paddle/fluid/tests/unittests/test_imperative.py
+++ b/python/paddle/fluid/tests/unittests/test_imperative.py
@@ -81,8 +81,8 @@ class TestImperative(unittest.TestCase):
     def test_pylayer(self):
         with fluid.imperative.guard():
             my_py_layer = MyPyLayer()
-            out = my_py_layer([np.ones([2, 2], np.float32)])
-            sys.stderr.write('%s\n' % np.array(out))
+            outs = my_py_layer([np.ones([2, 2], np.float32)])
+            sys.stderr.write('%s\n' % outs[0]._numpy())
             # out.backward()
 
     def test_layer_in_out(self):
-- 
GitLab