Imperative remove all descs (#16045)

* Remove Desc in Forward Pass

* Refactor VarBase

* Add dbg info

* Only check type in imperative mode

* Polish code and support optimizer

test=develop

* Fix stop gradient problem in PyLayer

test=develop

paddle/fluid/imperative/layer.cc

@@ -159,10 +159,9 @@ class Autograd {
       for (auto it : candidate->pre_ops_) {
         for (OpBase* pre_op : it.second) {
           if (!pre_op) continue;
-          VLOG(5) << "op dep " << candidate->op_desc_->Type() << " trace id "
+          VLOG(5) << "op dep " << candidate->Type() << " trace id "
                   << candidate->trace_id_ << " <---- " << it.first << " <---- "
-                  << pre_op->op_desc_->Type() << " trace id "
-                  << pre_op->trace_id_;
+                  << pre_op->Type() << " trace id " << pre_op->trace_id_;
           if (visited.find(pre_op) == visited.end()) {
             visited.insert(pre_op);
             queue.push_back(pre_op);
@@ -180,10 +179,12 @@ std::unique_ptr<VarBase> VarBase::NewVarBase(const platform::Place& dst_place,
   PADDLE_ENFORCE(var_->IsInitialized(),
                  "Variable must be initialized when getting numpy tensor");
 
-  std::unique_ptr<VarBase> new_var(new VarBase());
+  // TODO(minqiyang): change this after move unique_name generator to CXX
+  const framework::LoDTensor& self_tensor = var_->Get<framework::LoDTensor>();
+  std::unique_ptr<VarBase> new_var(new VarBase(
+      "Itmp", self_tensor.type(), self_tensor.dims(), dst_place, true, false));
   framework::LoDTensor* tensor =
       new_var->var_->GetMutable<framework::LoDTensor>();
-  tensor->Resize(var_->Get<framework::LoDTensor>().dims());
   tensor->set_lod(var_->Get<framework::LoDTensor>().lod());
 
   if (blocking) {
@@ -199,52 +200,62 @@ std::unique_ptr<VarBase> VarBase::NewVarBase(const platform::Place& dst_place,
   }
 
   if (platform::is_gpu_place(dst_place)) {
-    VLOG(3) << "copy tensor " << var_desc_->Name() << " from gpu";
+    VLOG(3) << "copy tensor " << Name() << " from gpu";
   }
 
   return new_var;
 }
 
 framework::LoDTensor& VarBase::GradValue() {
-  VLOG(3) << "get var grad " << var_desc_->Name();
+  VLOG(3) << "get var grad " << Name();
+  PADDLE_ENFORCE_NOT_NULL(grads_,
+                          "Could not get grad value from no grad variable");
   return *(grads_->var_->GetMutable<framework::LoDTensor>());
 }
 
 std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
   if (grad_op_descs_.empty() && backward_id_ <= 0) {
-    VLOG(3) << "op with no grad: " << op_desc_->Type();
+    VLOG(3) << "op with no grad: " << Type();
     return {};
   }
 
-  VLOG(3) << "apply op grad: " << op_desc_->Type();
-  std::vector<framework::VariableValueMap> grad_outputs;
+  VLOG(3) << "apply op grad: " << Type();
+  std::vector<framework::VariableValueMap> tmp_grad_outputs;
   if (backward_id_ > 0) {
     VLOG(3) << "py_layer_grad";
-    grad_outputs.resize(1);
-    grad_outputs[0][framework::GradVarName(PyLayer::kFwdOut)] =
+    tmp_grad_outputs.resize(1);
+    tmp_grad_outputs[0][framework::GradVarName(PyLayer::kFwdOut)] =
         PyLayer::ApplyGrad(
             backward_id_,
             grad_input_vars_[0][framework::GradVarName(PyLayer::kFwdInp)]);
   } else {
-    grad_outputs.resize(grad_op_descs_.size());
-    for (size_t k = 0; k < grad_op_descs_.size(); ++k) {
+    const size_t grad_op_count = grad_op_descs_.size();
+
+    tmp_grad_outputs.resize(grad_op_count);
+    for (size_t k = 0; k < grad_op_count; ++k) {
       framework::OpDesc* grad_op_desc = grad_op_descs_[k];
-      VLOG(3) << "op grad " << grad_op_desc->Type();
-      for (auto it : grad_output_vars_[k]) {
-        auto& outputs = grad_outputs[k][it.first];
+      auto& grad_output_variable_map = grad_output_vars_[k];
+
+      VLOG(3) << "apply grad op " << grad_op_desc->Type();
+
+      // Allocate tmp grad output variable
+      for (auto it : grad_output_variable_map) {
+        auto& outputs = tmp_grad_outputs[k][it.first];
+        outputs.reserve(it.second.size());
         for (size_t i = 0; i < it.second.size(); ++i) {
           // Allocate a new variable
           Variable* tmp_var = new framework::Variable();
           tmp_var->GetMutable<framework::LoDTensor>();
-          outputs.push_back(tmp_var);
+          outputs.emplace_back(tmp_var);
         }
       }
 
-      framework::RuntimeContext ctx(grad_input_vars_[k], grad_outputs[k]);
+      // Run grad op
+      framework::RuntimeContext ctx(grad_input_vars_[k], tmp_grad_outputs[k]);
 
       // No need to do compile time infer shape here.
       // grad_op_desc_->InferShape(*block_);
-      grad_op_desc->InferVarType(block_);
+      // grad_op_desc->InferVarType(block_);
 
       std::unique_ptr<framework::OperatorBase> opbase =
           framework::OpRegistry::CreateOp(*grad_op_desc);
@@ -260,9 +271,10 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
     }
   }
 
+  // Add tmp grad outputs to original grad vars
   for (size_t k = 0; k < grad_output_vars_.size(); ++k) {
     for (auto it : grad_output_vars_[k]) {
-      auto& outputs = grad_outputs[k][it.first];
+      auto& outputs = tmp_grad_outputs[k][it.first];
       auto& origin_outputs = it.second;
       PADDLE_ENFORCE_EQ(outputs.size(), origin_outputs.size());
 
@@ -316,19 +328,14 @@ void PyLayer::RegisterFunc(int func_id, const py::object& py_func) {
 
 int PyLayer::NumFuncs() { return py_funcs_.size(); }
 
-std::vector<VarBase*> PyLayer::Apply(int func_id,
+std::vector<Variable*> PyLayer::Apply(int func_id,
                                       const std::vector<VarBase*>& inputs) {
   std::vector<framework::Variable*> invars;
   for (const VarBase* in : inputs) {
     invars.push_back(in->var_);
   }
   PADDLE_ENFORCE(py_funcs_.find(func_id) != py_funcs_.end());
-  std::vector<Variable*> outvars = CallPythonFunc(py_funcs_[func_id], invars);
-  std::vector<VarBase*> ret;
-  for (Variable* v : outvars) {
-    ret.push_back(new VarBase(v, new VarBase(true)));
-  }
-  return ret;
+  return CallPythonFunc(py_funcs_[func_id], invars);
 }
 
 std::vector<Variable*> PyLayer::ApplyGrad(

paddle/fluid/imperative/layer.h

@@ -112,31 +112,53 @@ class OpBase;
  */
 class VarBase {
  public:
-  VarBase() : VarBase(new framework::Variable(), new VarBase(true)) {}
-
-  explicit VarBase(bool stop_gradient)
-      : VarBase(new framework::Variable(),
-                stop_gradient ? nullptr : new VarBase(true), stop_gradient) {}
-
-  VarBase(framework::Variable* var, VarBase* grad)
-      : VarBase(var, grad, false) {}
+  // Internal interface, create VarBase from exist variable
+  VarBase(const std::string& name, framework::Variable* var, VarBase* grad,
+          bool stop_gradient)
+      : VarBase(name, var->Get<framework::LoDTensor>().type(),
+                var->Get<framework::LoDTensor>().dims(),
+                var->Get<framework::LoDTensor>().place(), var, grad,
+                stop_gradient, false) {}
+
+  // Python interface
+  VarBase(const std::string& name, const framework::proto::VarType::Type dtype,
+          const std::vector<int64_t>& shape, const platform::Place& place,
+          bool stop_gradient, bool persistable)
+      : VarBase(name, dtype, framework::make_ddim(shape), place, stop_gradient,
+                persistable) {}
+
+  // Internal interface, create VarBase from with ddim
+  VarBase(const std::string& name, const framework::proto::VarType::Type dtype,
+          const framework::DDim& shape, const platform::Place& place,
+          bool stop_gradient, bool persistable)
+      : VarBase(name, dtype, shape, place, nullptr, nullptr, stop_gradient,
+                persistable) {}
 
  private:
-  VarBase(framework::Variable* var, VarBase* grad, bool stop_gradient)
-      : name_(),
-        var_desc_(nullptr),
+  VarBase(const std::string& name, framework::proto::VarType::Type dtype,
+          const framework::DDim& shape, const platform::Place& place,
+          framework::Variable* var, VarBase* grad, bool stop_gradient,
+          bool persistable)
+      : name_(name),
+        dtype_(dtype),
+        place_(place),
         var_(var),
         grads_(grad),
-        block_(nullptr),
-        persistable_(false),
         stop_gradient_(stop_gradient),
+        persistable_(persistable),
         pre_op_(nullptr),
         pre_op_out_name_(),
-        pre_op_out_idx_(-1) {}
+        pre_op_out_idx_(-1) {
+    if (!var_) {
+      var_ = new framework::Variable();
+      auto tensor = var_->GetMutable<framework::LoDTensor>();
+      tensor->Resize(shape);
+      tensor->mutable_data(place_, dtype_);
+    }
+  }
 
  public:
   virtual ~VarBase() {
-    // TODO(minqiyang): remove var desc from block desc
     if (var_) {
       delete var_;
       var_ = nullptr;
@@ -151,14 +173,30 @@ class VarBase {
     pre_op_out_idx_ = -1;
   }
 
-  inline OpBase* PreOp() const { return pre_op_; }
-  inline int PreOpOutIdx() const { return pre_op_out_idx_; }
+  inline void SetName(const std::string& name) { name_ = name; }
+  inline std::string Name() const { return name_; }
+
+  inline std::vector<int64_t> Shape() const {
+    if (var_->IsInitialized()) {
+      return framework::vectorize(var_->Get<framework::LoDTensor>().dims());
+    } else {
+      return {};
+    }
+  }
+
+  inline framework::proto::VarType::Type DType() const { return dtype_; }
 
   inline void SetStopGradient(bool stop_gradient) {
     stop_gradient_ = stop_gradient;
   }
   inline bool IsStopGradient() const { return stop_gradient_; }
 
+  inline void SetPersistable(bool persistable) { persistable_ = persistable; }
+  inline bool IsPersistable() const { return persistable_; }
+
+  inline OpBase* PreOp() const { return pre_op_; }
+  inline int PreOpOutIdx() const { return pre_op_out_idx_; }
+
   void RunBackward();
 
   inline void ResetPreOp(OpBase* op) {
@@ -180,7 +218,7 @@ class VarBase {
   }
 
   void ClearGradient() {
-    VLOG(1) << "clear gradient of " << var_desc_->Name();
+    VLOG(1) << "clear gradient of " << Name();
     if (grads_ && grads_->var_ && grads_->var_->IsInitialized()) {
       auto grads_t = grads_->var_->GetMutable<framework::LoDTensor>();
       operators::math::set_constant(
@@ -196,23 +234,20 @@ class VarBase {
                                       const bool blocking) const;
 
   inline std::string GradName() const {
-    PADDLE_ENFORCE(
-        var_desc_,
-        "Couldn't get gradient variable's name, please call backward() first");
-    return string::Sprintf("%s@IGrad", var_desc_->Name());
+    return string::Sprintf("%s@IGrad", Name());
   }
 
   std::string name_;
-  framework::VarDesc* var_desc_;
+  framework::proto::VarType::Type dtype_;
+  platform::Place place_;
 
   framework::Variable* var_;
   VarBase* grads_;
 
-  framework::BlockDesc* block_;
-  bool persistable_;
-
  private:
   bool stop_gradient_;
+  bool persistable_;
+
   OpBase* pre_op_;
   std::string pre_op_out_name_;
   int pre_op_out_idx_;
@@ -223,11 +258,11 @@ class VarBase {
  */
 class PYBIND11_HIDDEN OpBase {
  public:
-  OpBase()
-      : op_desc_(nullptr),
+  OpBase(const std::string& type)
+      : type_(type),
+        trace_id_(-1),
         forward_id_(-1),
         backward_id_(-1),
-        trace_id_(-1),
         place_(platform::CPUPlace()),
         backward_hooks_() {}
 
@@ -249,13 +284,34 @@ class PYBIND11_HIDDEN OpBase {
 
   std::map<std::string, std::vector<VarBase*>> ApplyGrad();
 
+  inline std::string Type() const { return type_; }
+  inline std::string GradOpType(size_t index) const {
+    PADDLE_ENFORCE_NOT_NULL(grad_op_descs_[index]);
+    return grad_op_descs_[index]->Type();
+  }
+
   void RegisterBackwardHooks(const py::object& callable);
 
   void InvokeBackwardHooks();
 
-  // One of `op_desc_` or `forward_id_` is set, not both.
-  // For pure python PyLayer, use `forward_id_`, otherwise, use op_desc_.
-  framework::OpDesc* op_desc_;
+  void TrackPreOp(const VarBase* inp_var, const std::string& inp_name) {
+    if (inp_var->PreOp() && !inp_var->IsStopGradient()) {
+      VLOG(3) << "add pre op " << inp_var->PreOp()->Type() << " in slot "
+              << inp_name;
+      pre_ops_[inp_name].push_back(inp_var->PreOp());
+      pre_ops_out_idx_[inp_name].push_back(inp_var->PreOpOutIdx());
+    } else {
+      VLOG(3) << "no pre op in slot " << inp_name
+              << " input var stop_gradient: " << inp_var->IsStopGradient();
+      pre_ops_[inp_name].push_back(nullptr);
+      // pre_ops_out_idx_[inp_name].push_back(-1);
+    }
+  }
+
+  std::string type_;
+  // One of `trace_id_` or `forward_id_` is set, not both.
+  // For pure python PyLayer, use `forward_id_`, otherwise, use trace_id_.
+  int trace_id_;
   int forward_id_;
 
   // When has backward, one of `grad_op_descs_` or `backward_id_` is set,
@@ -263,7 +319,6 @@ class PYBIND11_HIDDEN OpBase {
   // Note: each fwd op corresponds to a vector of bwd ops.
   std::vector<framework::OpDesc*> grad_op_descs_;
   int backward_id_;
-  int trace_id_;
 
   platform::Place place_;
 
@@ -277,8 +332,6 @@ class PYBIND11_HIDDEN OpBase {
   // Outputs to a vector of bwd ops.
   std::vector<framework::VariableValueMap> grad_output_vars_;
 
-  framework::BlockDesc* block_;
-
   std::vector<py::object> backward_hooks_;
 };
 
@@ -303,8 +356,8 @@ class PyLayer {
 
   static int NumFuncs();
 
-  static std::vector<VarBase*> Apply(int func_id,
-                                     const std::vector<VarBase*>& inputs);
+  static std::vector<framework::Variable*> Apply(
+      int func_id, const std::vector<VarBase*>& inputs);
 
   static std::vector<framework::Variable*> ApplyGrad(
       int func_id, const std::vector<framework::Variable*>& inputs);

paddle/fluid/imperative/tracer.cc

@@ -56,15 +56,19 @@ void CreateGradOp(const framework::OpDesc& op_desc,
   }
 }
 
-void InitVar(framework::Variable* var, framework::Variable* grad_var,
-             platform::DeviceContext* dev_ctx) {
+void InitGrad(VarBase* var, platform::DeviceContext* dev_ctx) {
+  PADDLE_ENFORCE_NOT_NULL(var, "Could not get valid var base");
   PADDLE_ENFORCE_NOT_NULL(dev_ctx,
                           "Could not get valid device from forward op");
-  auto& var_t = var->Get<framework::LoDTensor>();
-  grad_var->GetMutable<framework::LoDTensor>()->mutable_data<float>(
-      var_t.dims(), dev_ctx->GetPlace());
-  operators::math::set_constant(
-      *dev_ctx, grad_var->GetMutable<framework::LoDTensor>(), 0.0);
+
+  if (var->grads_ == nullptr) {
+    auto& var_t = var->var_->Get<framework::LoDTensor>();
+    var->grads_ = new VarBase(var->GradName(), framework::proto::VarType::FP32,
+                              framework::vectorize(var_t.dims()),
+                              dev_ctx->GetPlace(), true, false);
+    auto grad_t = var->grads_->var_->GetMutable<framework::LoDTensor>();
+    operators::math::set_constant(*dev_ctx, grad_t, 0.0);
+  }
 }
 
 platform::Place GetExpectedPlace(platform::Place place, VarBasePtrMap inputs) {
@@ -85,6 +89,62 @@ platform::Place GetExpectedPlace(platform::Place place, VarBasePtrMap inputs) {
   return result;
 }
 
+framework::VariableNameMap CreateInputVarNameMap(
+    const OpBase* op, const VarBasePtrMap& varbase_map) {
+  framework::VariableNameMap result;
+
+  auto& info_map = framework::OpInfoMap::Instance();
+  auto* op_info = info_map.GetNullable(op->Type());
+  if (op_info == nullptr || op_info->proto_ == nullptr) {
+    return result;
+  }
+
+  for (auto& in : op_info->Proto().inputs()) {
+    auto it = varbase_map.find(in.name());
+    if (it == varbase_map.end()) {
+      PADDLE_ENFORCE(in.dispensable());
+      result[in.name()] = {};
+    } else {
+      auto var_vector = it->second;
+      std::vector<std::string> args;
+      args.reserve(var_vector.size());
+      for (VarBase* var_base : var_vector) {
+        args.emplace_back(var_base->Name());
+      }
+      result[in.name()] = args;
+    }
+  }
+  return result;
+}
+
+framework::VariableNameMap CreateOutputVarNameMap(
+    const OpBase* op, const VarBasePtrMap& varbase_map) {
+  framework::VariableNameMap result;
+
+  auto& info_map = framework::OpInfoMap::Instance();
+  auto* op_info = info_map.GetNullable(op->Type());
+  if (op_info == nullptr || op_info->proto_ == nullptr) {
+    return result;
+  }
+
+  for (auto& out : op_info->Proto().outputs()) {
+    auto it = varbase_map.find(out.name());
+    if (it == varbase_map.end()) {
+      PADDLE_ENFORCE(out.dispensable());
+      result[out.name()] = {};
+    } else {
+      auto var_vector = it->second;
+      std::vector<std::string> args;
+      args.reserve(var_vector.size());
+      for (VarBase* var_base : var_vector) {
+        args.emplace_back(var_base->Name());
+      }
+      result[out.name()] = args;
+    }
+  }
+  return result;
+}
+
 Tracer::Tracer(framework::BlockDesc* root_block) : root_block_(root_block) {
   if (!FLAGS_tracer_profile_fname.empty()) {
     std::call_once(gTracerProfileOnce, [] {
@@ -101,7 +161,7 @@ Tracer::Tracer(framework::BlockDesc* root_block) : root_block_(root_block) {
 
 std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
                                     const VarBasePtrMap& outputs,
-                                    framework::BlockDesc* block,
+                                    framework::AttributeMap attrs_map,
                                     const platform::Place expected_place,
                                     const bool stop_gradient) {
 #ifdef WITH_GPERFTOOLS
@@ -110,40 +170,27 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
   }
 #endif
 
-  std::map<std::string, VarBase*> vars;
-
-  framework::OpDesc* op_desc = op->op_desc_;
-  VLOG(3) << "tracer tracing " << op_desc->Type() << " trace id "
-          << op->trace_id_;
-  op_desc->InferShape(*block);
-  op_desc->InferVarType(block);
-
-  std::unique_ptr<framework::OperatorBase> op_base =
-      framework::OpRegistry::CreateOp(*op_desc);
-
   framework::VariableValueMap invars_map;
   framework::VariableValueMap outvars_map;
 
+  // Construct input_vars_map and output_vars_map
+  std::map<std::string, VarBase*> current_vars_map;
   op->input_vars_ = inputs;
   for (auto it : op->input_vars_) {
     auto& invars = invars_map[it.first];
     invars.reserve(it.second.size());
     for (VarBase* inp : it.second) {
-      PADDLE_ENFORCE_NOT_NULL(inp->var_, "op %s input %s nullptr",
-                              op->op_desc_->Type(), inp->var_desc_->Name());
+      PADDLE_ENFORCE_NOT_NULL(inp->var_, "op %s input %s nullptr", op->Type(),
+                              inp->Name());
 
       invars.emplace_back(inp->var_);
-      vars[inp->var_desc_->Name()] = inp;
-      if (inp->PreOp() && !inp->IsStopGradient()) {
-        op->pre_ops_[it.first].push_back(inp->PreOp());
-        op->pre_ops_out_idx_[it.first].push_back(inp->PreOpOutIdx());
-        VLOG(3) << "add pre op " << inp->PreOp()->op_desc_->Type();
-      } else {
-        op->pre_ops_[it.first].push_back(nullptr);
+      op->TrackPreOp(inp, it.first);
+      if (!stop_gradient) {
+        current_vars_map[inp->Name()] = inp;
       }
-      VLOG(3) << "input vname " << inp->var_desc_->Name() << " "
-              << inp->var_->IsInitialized() << " stop_gradient "
-              << inp->IsStopGradient();
+      VLOG(3) << "input var name: " << inp->Name()
+              << " inited: " << inp->var_->IsInitialized()
+              << " stop_grad: " << inp->IsStopGradient();
     }
   }
 
@@ -152,25 +199,38 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
     auto& outvars = outvars_map[it.first];
     const std::vector<VarBase*>& outputs = it.second;
     outvars.reserve(outputs.size());
-    for (size_t i = 0; i < outputs.size(); ++i) {
+    for (size_t i = 0U; i < outputs.size(); ++i) {
       VarBase* out = outputs[i];
       outvars.emplace_back(out->var_);
-      vars[out->var_desc_->Name()] = out;
-
-      framework::VarDesc* var_desc = block->FindVar(out->var_desc_->Name());
-      if (var_desc->GetType() == framework::proto::VarType::LOD_TENSOR) {
-        out->var_->GetMutable<framework::LoDTensor>();
-      } else {
-        LOG(ERROR) << "tracer doesn't support yet";
-      }
       out->TrackPreOp(op, it.first, i, stop_gradient);
+      if (!stop_gradient) {
+        current_vars_map[out->Name()] = out;
+      }
 
-      VLOG(3) << "output vname " << out->var_desc_->Name() << " "
-              << out->var_->IsInitialized();
+      VLOG(3) << "input var name: " << out->Name()
+              << " inited: " << out->var_->IsInitialized()
+              << " stop_grad: " << out->IsStopGradient();
+    }
   }
+
+  // Check attrs and create op
+  framework::VariableNameMap invars_name_map =
+      CreateInputVarNameMap(op, inputs);
+  framework::VariableNameMap outvars_name_map =
+      CreateOutputVarNameMap(op, outputs);
+
+  auto& info = framework::OpInfoMap::Instance().Get(op->Type());
+  if (info.Checker() != nullptr) {
+    info.Checker()->Check(&attrs_map);
   }
 
-  VLOG(3) << "tracer running " << op_desc->Type();
+  std::unique_ptr<framework::OperatorBase> op_base =
+      framework::OpRegistry::CreateOp(op->Type(), invars_name_map,
+                                      outvars_name_map, attrs_map);
+
+  // TODO(minqiyang): Support infer var type in imperative mode
+  // Run forward op
+  VLOG(3) << "tracer running " << op->Type();
   framework::RuntimeContext ctx(invars_map, outvars_map);
 
   // TODO(panyx0718): Cache p.
@@ -186,36 +246,44 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
       framework::ExecutionContext(prepared_op.op, scope, *prepared_op.dev_ctx,
                                   prepared_op.ctx, prepared_op.kernel_configs));
 
+  // construct backward op
   std::set<std::string> vars_saved_for_backward;
-
   if (!stop_gradient) {
+    VLOG(5) << "start construct backward op";
+
+    // construct grad op descs
+    std::unique_ptr<framework::OpDesc> fwd_op_desc(new framework::OpDesc(
+        op->Type(), invars_name_map, outvars_name_map, attrs_map));
     std::unique_ptr<std::unordered_map<std::string, std::string>> grad_to_var(
         new std::unordered_map<std::string, std::string>());
-    CreateGradOp(*op_desc, {}, {block}, &op->grad_op_descs_, grad_to_var.get());
+    // NOTE(minqiyang): We don't support control flow op in imperative now
+    // Add grad_block_ when we want to support it
+    CreateGradOp(*fwd_op_desc, {}, {}, &op->grad_op_descs_, grad_to_var.get());
 
-    op->grad_input_vars_.resize(op->grad_op_descs_.size());
-    op->grad_output_vars_.resize(op->grad_op_descs_.size());
+    VLOG(5) << "create grad op desc: " << op->grad_op_descs_[0]->Type();
 
-    for (size_t i = 0; i < op->grad_op_descs_.size(); ++i) {
+    const size_t grad_op_count = op->grad_op_descs_.size();
+
+    op->grad_input_vars_.resize(grad_op_count);
+    op->grad_output_vars_.resize(grad_op_count);
+
+    for (size_t i = 0; i < grad_op_count; ++i) {
       framework::OpDesc* grad_op_desc = op->grad_op_descs_[i];
       for (auto it : grad_op_desc->Inputs()) {
         auto& grad_in_vars = op->grad_input_vars_[i][it.first];
+        grad_in_vars.reserve(it.second.size());
         for (const std::string& grad_invar : it.second) {
-          block->FindRecursiveOrCreateVar(grad_invar);
           auto var_it = grad_to_var->find(grad_invar);
           if (var_it == grad_to_var->end()) {
-            auto fwd_var_it = vars.find(grad_invar);
-            PADDLE_ENFORCE(fwd_var_it != vars.end());
+            auto fwd_var_it = current_vars_map.find(grad_invar);
+            PADDLE_ENFORCE(fwd_var_it != current_vars_map.end());
             // Forward inputs or outputs.
-            grad_in_vars.push_back(fwd_var_it->second->var_);
+            grad_in_vars.emplace_back(fwd_var_it->second->var_);
           } else {
-            VarBase* var = vars[var_it->second];
-            if (!var->grads_->var_->IsInitialized()) {
-              InitVar(var->var_, var->grads_->var_,
-                      prepared_op.GetDeviceContext());
-            }
+            VarBase* var = current_vars_map[var_it->second];
+            InitGrad(var, prepared_op.GetDeviceContext());
             // Douts.
-            grad_in_vars.push_back(var->grads_->var_);
+            grad_in_vars.emplace_back(var->grads_->var_);
           }
 
           vars_saved_for_backward.insert(it.first);
@@ -225,48 +293,48 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
       for (auto it : grad_op_desc->Outputs()) {
         auto& grad_out_vars = op->grad_output_vars_[i][it.first];
         for (const std::string& grad_outvar : it.second) {
-          block->FindRecursiveOrCreateVar(grad_outvar);
           auto var_it = grad_to_var->find(grad_outvar);
           PADDLE_ENFORCE(var_it != grad_to_var->end(),
                          "Could not found the grad op output var, should this "
                          "operator %s's stop gradient be True",
-                         op_desc->Type());
-          VarBase* var = vars[var_it->second];
-          if (!var->grads_->var_->IsInitialized()) {
-            InitVar(var->var_, var->grads_->var_,
-                    prepared_op.GetDeviceContext());
-          }
+                         op->Type());
+          VarBase* var = current_vars_map[var_it->second];
+          InitGrad(var, prepared_op.GetDeviceContext());
           grad_out_vars.push_back(var->grads_->var_);
         }
       }
     }
   }
 
-  op->block_ = block;
   return vars_saved_for_backward;
 }
 
 std::vector<VarBase*> Tracer::PyTrace(OpBase* op,
                                       const std::vector<VarBase*>& inputs,
                                       bool stop_gradient) {
-  VLOG(3) << "py_trace";
+  VLOG(3) << "py_trace " << op->Type();
+
   op->input_vars_[PyLayer::kFwdInp] = inputs;
-  op->output_vars_[PyLayer::kFwdOut] = PyLayer::Apply(op->forward_id_, inputs);
+
+  std::vector<framework::Variable*> ret_vars =
+      PyLayer::Apply(op->forward_id_, inputs);
+
   for (VarBase* inp : inputs) {
-    if (inp->PreOp() && !inp->IsStopGradient()) {
-      op->pre_ops_[PyLayer::kFwdInp].push_back(inp->PreOp());
-      op->pre_ops_out_idx_[PyLayer::kFwdInp].push_back(inp->PreOpOutIdx());
-    } else {
-      op->pre_ops_[PyLayer::kFwdInp].push_back(nullptr);
-    }
+    op->TrackPreOp(inp, PyLayer::kFwdInp);
   }
 
-  auto& outputs = op->output_vars_[PyLayer::kFwdOut];
-  for (size_t i = 0; i < outputs.size(); ++i) {
-    VarBase* out = outputs[i];
+  std::vector<VarBase*>& outputs = op->output_vars_[PyLayer::kFwdOut];
+  outputs.reserve(ret_vars.size());
+  for (size_t i = 0U; i != ret_vars.size(); ++i) {
+    framework::Variable* v = ret_vars[i];
+    VarBase* out = new VarBase(string::Sprintf("%s_out_%d", op->Type(), i), v,
+                               nullptr, stop_gradient);
+    outputs.emplace_back(out);
     out->TrackPreOp(op, PyLayer::kFwdOut, i, stop_gradient);
   }
+
   if (!stop_gradient) {
+    VLOG(5) << "start construct backward op";
     op->grad_input_vars_.resize(1);
     op->grad_output_vars_.resize(1);
     auto& grad_input_vars =
@@ -281,23 +349,16 @@ std::vector<VarBase*> Tracer::PyTrace(OpBase* op,
       grad_input_vars.push_back(out->var_);
     }
 
+    // TODO(minqiyang): Add GPU support for PyLayer, only support CPU now
     platform::CPUPlace place;
     for (VarBase* out : outputs) {
+      InitGrad(out, platform::DeviceContextPool::Instance().Get(place));
       grad_input_vars.push_back(out->grads_->var_);
-      if (!grad_input_vars.back()->IsInitialized()) {
-        // TODO(minqiyang): Add GPU support for PyLayer, only support CPU now
-        InitVar(out->var_, grad_input_vars.back(),
-                platform::DeviceContextPool::Instance().Get(place));
-      }
     }
 
-    for (const VarBase* inp : inputs) {
+    for (VarBase* inp : inputs) {
+      InitGrad(inp, platform::DeviceContextPool::Instance().Get(place));
       grad_output_vars.push_back(inp->grads_->var_);
-      if (!grad_output_vars.back()->IsInitialized()) {
-        // TODO(minqiyang): Add GPU support for PyLayer, only support CPU now
-        InitVar(inp->var_, grad_output_vars.back(),
-                platform::DeviceContextPool::Instance().Get(place));
-      }
     }
   }
   return outputs;

paddle/fluid/imperative/tracer.h

@@ -17,6 +17,8 @@
 #include <map>
 #include <set>
 #include <string>
+#include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "paddle/fluid/framework/op_desc.h"
@@ -34,7 +36,8 @@ void CreateGradOp(const framework::OpDesc& op_desc,
                   framework::OpDesc** grad_op_desc,
                   std::unordered_map<std::string, std::string>* grad_to_var);
 
-void InitVar(framework::Variable* var, framework::Variable* grad_var);
+void InitVar(const VarBase* var, framework::Variable* grad_var,
+             platform::DeviceContext* dev_ctx);
 
 platform::Place GetExpectedPlace(platform::Place place, VarBasePtrMap inputs);
 
@@ -46,7 +49,7 @@ class Tracer {
 
   std::set<std::string> Trace(OpBase* op, const VarBasePtrMap& inputs,
                               const VarBasePtrMap& outputs,
-                              framework::BlockDesc* block,
+                              framework::AttributeMap attrs_map,
                               const platform::Place expected_place,
                               const bool stop_gradient = false);
 

paddle/fluid/memory/allocation/legacy_allocator.cc

@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/memory/allocation/legacy_allocator.h"
 
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>

paddle/fluid/pybind/imperative.cc

@@ -13,10 +13,18 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/pybind/imperative.h"
+
+#include <pybind11/chrono.h>
+#include <pybind11/complex.h>
+#include <pybind11/functional.h>
+#include <pybind11/stl.h>
+
 #include "paddle/fluid/framework/block_desc.h"
 #include "paddle/fluid/imperative/tracer.h"
 #include "paddle/fluid/imperative/type_defs.h"
 
+#include "paddle/fluid/pybind/pybind_boost_headers.h"
+
 namespace paddle {
 namespace pybind {
 
@@ -31,20 +39,20 @@ void BindTracer(pybind11::module* m) {
            [](imperative::Tracer& self, imperative::OpBase* op,
               const imperative::VarBasePtrMap& inputs,
               const imperative::VarBasePtrMap& outputs,
-              framework::BlockDesc* block,
+              framework::AttributeMap attrs_map,
               const platform::CPUPlace expected_place,
               const bool stop_gradient = false) {
-             return self.Trace(op, inputs, outputs, block, expected_place,
+             return self.Trace(op, inputs, outputs, attrs_map, expected_place,
                                stop_gradient);
            })
       .def("trace",
            [](imperative::Tracer& self, imperative::OpBase* op,
               const imperative::VarBasePtrMap& inputs,
               const imperative::VarBasePtrMap& outputs,
-              framework::BlockDesc* block,
+              framework::AttributeMap attrs_map,
               const platform::CUDAPlace expected_place,
               const bool stop_gradient = false) {
-             return self.Trace(op, inputs, outputs, block, expected_place,
+             return self.Trace(op, inputs, outputs, attrs_map, expected_place,
                                stop_gradient);
            })
       .def("py_trace", &imperative::Tracer::PyTrace,

paddle/fluid/pybind/imperative.h

@@ -14,6 +14,7 @@ limitations under the License. */
 #pragma once
 
 #include <Python.h>
+#include <string>
 #include <vector>
 #include "paddle/fluid/imperative/layer.h"
 #include "pybind11/pybind11.h"
@@ -36,6 +37,8 @@ class Layer : public imperative::Layer {
 class PYBIND11_HIDDEN PyOpBase : public imperative::OpBase {
  public:
   using imperative::OpBase::OpBase;  // Inherit constructors
+
+  PyOpBase(const std::string& name) : OpBase(name) {}
 };
 
 class PyVarBase : public imperative::VarBase {

paddle/fluid/pybind/protobuf.cc

@@ -23,97 +23,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/framework/var_desc.h"
 
-// Cast boost::variant for PyBind.
-// Copy from
-// https://github.com/pybind/pybind11/issues/576#issuecomment-269563199
-namespace pybind11 {
-namespace detail {
-
-#if !defined(PYBIND11_HIDDEN)
-#ifdef _WIN32
-#define PYBIND11_HIDDEN __declspec(dllexport)
-#else
-#define PYBIND11_HIDDEN __attribute__((visibility("hidden")))
-#endif
-#endif
-
-// Can be replaced by a generic lambda in C++14
-struct PYBIND11_HIDDEN paddle_variant_caster_visitor
-    : public boost::static_visitor<handle> {
-  return_value_policy policy;
-  handle parent;
-
-  paddle_variant_caster_visitor(return_value_policy policy, handle parent)
-      : policy(policy), parent(parent) {}
-
-  template <class T>
-  handle operator()(T const &src) const {
-    return make_caster<T>::cast(src, policy, parent);
-  }
-};
-
-template <class Variant>
-struct paddle_variant_caster;
-
-template <template <class...> class V, class... Ts>
-struct paddle_variant_caster<V<Ts...>> {
-  using Type = V<Ts...>;
-
-  template <typename T>
-  typename std::enable_if<
-      !std::is_same<T, boost::detail::variant::void_>::value, bool>::type
-  try_load(handle src, bool convert) {
-    auto caster = make_caster<T>();
-    if (!load_success_ && caster.load(src, convert)) {
-      load_success_ = true;
-
-      if (std::is_same<T, std::vector<float>>::value) {
-        auto caster_ints = make_caster<std::vector<int64_t>>();
-        if (caster_ints.load(src, convert)) {
-          VLOG(4) << "This value are floats and int64_ts satisfy "
-                     "simultaneously, will set it's type to "
-                     "std::vector<int64_t>";
-          value = cast_op<std::vector<int64_t>>(caster_ints);
-          return true;
-        }
-      }
-
-      value = cast_op<T>(caster);
-      return true;
-    }
-    return false;
-  }
-
-  template <typename T>
-  typename std::enable_if<std::is_same<T, boost::detail::variant::void_>::value,
-                          bool>::type
-  try_load(handle src, bool convert) {
-    return false;
-  }
-
-  bool load(handle src, bool convert) {
-    auto unused = {false, try_load<Ts>(src, convert)...};
-    (void)(unused);
-    return load_success_;
-  }
-
-  static handle cast(Type const &src, return_value_policy policy,
-                     handle parent) {
-    paddle_variant_caster_visitor visitor(policy, parent);
-    return boost::apply_visitor(visitor, src);
-  }
-
-  PYBIND11_TYPE_CASTER(Type, _("Variant"));
-  bool load_success_{false};
-};
-
-// Add specialization for concrete variant type
-template <class... Args>
-struct type_caster<boost::variant<Args...>>
-    : paddle_variant_caster<boost::variant<Args...>> {};
-
-}  // namespace detail
-}  // namespace pybind11
+#include "paddle/fluid/pybind/pybind_boost_headers.h"
 
 namespace paddle {
 namespace pybind {

paddle/fluid/pybind/pybind.cc

@@ -149,8 +149,14 @@ PYBIND11_MODULE(core, m) {
         []() { return memory::allocation::GPUMemMonitor.PrintMemUsage(); });
 
   py::class_<imperative::VarBase>(m, "VarBase", R"DOC()DOC")
-      // .def(py::init<>())
-      .def(py::init<bool>(), py::arg("stop_gradient") = false)
+      .def(
+          py::init<const std::string &, paddle::framework::proto::VarType::Type,
+                   const std::vector<int64_t>, const paddle::platform::CPUPlace,
+                   bool, bool>())
+      .def(
+          py::init<const std::string &, paddle::framework::proto::VarType::Type,
+                   const std::vector<int64_t>,
+                   const paddle::platform::CUDAPlace, bool, bool>())
       .def("_run_backward",
            [](imperative::VarBase &self) { self.RunBackward(); })
       .def("_grad_name", &imperative::VarBase::GradName)
@@ -177,51 +183,21 @@ PYBIND11_MODULE(core, m) {
            py::return_value_policy::take_ownership)
       .def("value", [](const imperative::VarBase &self) { return self.var_; },
            py::return_value_policy::reference)
-      .def_property("name",
-                    [](const imperative::VarBase &self) { return self.name_; },
-                    [](imperative::VarBase &self, const std::string &name) {
-                      self.name_ = name;
-                    })
-      .def_property("block",
-                    [](const imperative::VarBase &self) { return self.block_; },
-                    [](imperative::VarBase &self, framework::BlockDesc *block) {
-                      self.block_ = block;
-                    },
-                    py::return_value_policy::reference)
-      .def_property(
-          "persistable",
-          [](const imperative::VarBase &self) { return self.persistable_; },
-          [](imperative::VarBase &self, const bool persistable) {
-            self.persistable_ = persistable;
-          })
-      .def_property(
-          "desc",
-          [](const imperative::VarBase &self) { return self.var_desc_; },
-          [](imperative::VarBase &self, framework::VarDesc *var_desc) {
-            self.var_desc_ = var_desc;
-          },
-          py::return_value_policy::reference)
-      .def_property(
-          "stop_gradient",
-          [](const imperative::VarBase &self) { return self.IsStopGradient(); },
-          [](imperative::VarBase &self, bool stop_gradient) {
-            self.SetStopGradient(stop_gradient);
-          });
+      .def_property("name", &imperative::VarBase::Name,
+                    &imperative::VarBase::SetName)
+      .def_property_readonly("shape", &imperative::VarBase::Shape)
+      .def_property_readonly("dtype", &imperative::VarBase::DType)
+      .def_property("persistable", &imperative::VarBase::IsPersistable,
+                    &imperative::VarBase::SetPersistable)
+      .def_property("stop_gradient", &imperative::VarBase::IsStopGradient,
+                    &imperative::VarBase::SetStopGradient);
 
   py::class_<imperative::OpBase, PyOpBase>(m, "OpBase", R"DOC()DOC")
-      .def(py::init<>())
+      .def(py::init<const std::string &>())
       .def("register_backward_hooks",
            [](imperative::OpBase &self, const py::object &callable) {
              self.RegisterBackwardHooks(callable);
            })
-      .def_property(
-          "desc", [](const imperative::OpBase &self) { return self.op_desc_; },
-          [](imperative::OpBase &self, framework::OpDesc *op_desc) {
-            if (op_desc) {
-              self.op_desc_ = op_desc;
-            }
-          },
-          py::return_value_policy::reference)
       .def_property("_trace_id",
                     [](const imperative::OpBase &self) {
                       pybind11::gil_scoped_release release;
@@ -260,7 +236,17 @@ PYBIND11_MODULE(core, m) {
           "apply",
           [](int func_id, const std::vector<imperative::VarBase *> &inputs)
               -> std::vector<imperative::VarBase *> {
-                return imperative::PyLayer::Apply(func_id, inputs);
+                auto ret_vars = imperative::PyLayer::Apply(func_id, inputs);
+                std::vector<imperative::VarBase *> outputs;
+                outputs.reserve(ret_vars.size());
+                for (size_t i = 0U; i != ret_vars.size(); ++i) {
+                  framework::Variable *v = ret_vars[i];
+                  // TODO(minqiyang): use unique_name generator to set a name
+                  outputs.emplace_back(
+                      new imperative::VarBase("", v, nullptr, true));
+                }
+
+                return outputs;
               },
           py::return_value_policy::take_ownership)
       .def_static("register_func",

paddle/fluid/pybind/pybind_boost_headers.h

+/* Copyright (c) 2018 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 <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/platform/variant.h"
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/stl.h"
+// Cast boost::variant for PyBind.
+// Copy from
+// https://github.com/pybind/pybind11/issues/576#issuecomment-269563199
+namespace pybind11 {
+namespace detail {
+
+#if !defined(PYBIND11_HIDDEN)
+#ifdef _WIN32
+#define PYBIND11_HIDDEN __declspec(dllexport)
+#else
+#define PYBIND11_HIDDEN __attribute__((visibility("hidden")))
+#endif
+#endif
+
+// Can be replaced by a generic lambda in C++14
+struct PYBIND11_HIDDEN paddle_variant_caster_visitor
+    : public boost::static_visitor<handle> {
+  return_value_policy policy;
+  handle parent;
+
+  paddle_variant_caster_visitor(return_value_policy policy, handle parent)
+      : policy(policy), parent(parent) {}
+
+  template <class T>
+  handle operator()(T const &src) const {
+    return make_caster<T>::cast(src, policy, parent);
+  }
+};
+
+template <class Variant>
+struct paddle_variant_caster;
+
+template <template <class...> class V, class... Ts>
+struct paddle_variant_caster<V<Ts...>> {
+  using Type = V<Ts...>;
+
+  template <typename T>
+  typename std::enable_if<
+      !std::is_same<T, boost::detail::variant::void_>::value, bool>::type
+  try_load(handle src, bool convert) {
+    auto caster = make_caster<T>();
+    if (!load_success_ && caster.load(src, convert)) {
+      load_success_ = true;
+
+      if (std::is_same<T, std::vector<float>>::value) {
+        auto caster_ints = make_caster<std::vector<int64_t>>();
+        if (caster_ints.load(src, convert)) {
+          VLOG(4) << "This value are floats and int64_ts satisfy "
+                     "simultaneously, will set it's type to "
+                     "std::vector<int64_t>";
+          value = cast_op<std::vector<int64_t>>(caster_ints);
+          return true;
+        }
+      }
+
+      value = cast_op<T>(caster);
+      return true;
+    }
+    return false;
+  }
+
+  template <typename T>
+  typename std::enable_if<std::is_same<T, boost::detail::variant::void_>::value,
+                          bool>::type
+  try_load(handle src, bool convert) {
+    return false;
+  }
+
+  bool load(handle src, bool convert) {
+    auto unused = {false, try_load<Ts>(src, convert)...};
+    (void)(unused);
+    return load_success_;
+  }
+
+  static handle cast(Type const &src, return_value_policy policy,
+                     handle parent) {
+    paddle_variant_caster_visitor visitor(policy, parent);
+    return boost::apply_visitor(visitor, src);
+  }
+
+  PYBIND11_TYPE_CASTER(Type, _("Variant"));
+  bool load_success_{false};
+};
+
+// Add specialization for concrete variant type
+template <class... Args>
+struct type_caster<boost::variant<Args...>>
+    : paddle_variant_caster<boost::variant<Args...>> {};
+
+}  // namespace detail
+}  // namespace pybind11

python/paddle/fluid/framework.py

@@ -304,10 +304,27 @@ class Variable(object):
                  is_data=False,
                  **kwargs):
         self.block = block
-        self.error_clip = error_clip
-
         if name is None:
             name = unique_name.generate('_generated_var')
+
+        if dtype is not None:
+            if not isinstance(dtype, core.VarDesc.VarType):
+                dtype = convert_np_dtype_to_dtype_(dtype)
+
+        if _in_imperative_mode():
+            # record vars in tracer rather than blocks
+            self._ivar = kwargs.get("ivar", None)
+            if not self._ivar:
+                self._ivar = core.VarBase(
+                    name, dtype if dtype else core.VarDesc.VarType.FP32,
+                    list(shape) if shape else [],
+                    _current_expected_place(), True
+                    if persistable else False, stop_gradient)
+            if persistable:
+                _imperative_tracer().trace_var(name, self)
+        else:
+            self.error_clip = error_clip
+
             is_new_var = False
             name = cpt.to_text(name)
             self.desc = self.block.desc.find_var(cpt.to_bytes(name))
@@ -319,10 +336,11 @@ class Variable(object):
             if is_new_var:
                 self.desc.set_type(type)
             elif self.desc.type() != type:
-            raise ValueError("Variable {0} has been created before. The "
+                raise ValueError(
+                    "Variable {0} has been created before. The "
                     "previous type is {1}; the new type is {2}. They"
-                             " are not matched".format(self.name,
-                                                       self.desc.type(), type))
+                    " are not matched".format(self.name, self.desc.type(),
+                                              type))
 
             if shape is not None:
                 if is_new_var:
@@ -336,25 +354,24 @@ class Variable(object):
                             "shape is {1}; the new shape is {2}. They are not "
                             "matched.".format(self.name, old_shape, shape))
             if dtype is not None:
-            if not isinstance(dtype, core.VarDesc.VarType):
-                dtype = convert_np_dtype_to_dtype_(dtype)
                 if is_new_var:
                     self.desc.set_dtype(dtype)
                 else:
                     old_dtype = self.dtype
                     if dtype != old_dtype:
-                    raise ValueError("Variable {0} has been created before. "
+                        raise ValueError(
+                            "Variable {0} has been created before. "
                             "The previous data type is {1}; the new "
                             "data type is {2}. They are not "
-                                     "matched.".format(self.name, old_dtype,
-                                                       dtype))
+                            "matched.".format(self.name, old_dtype, dtype))
 
             if lod_level is not None:
                 if is_new_var:
                     self.desc.set_lod_level(lod_level)
                 else:
                     if lod_level != self.lod_level:
-                    raise ValueError("Variable {0} has been created before. "
+                        raise ValueError(
+                            "Variable {0} has been created before. "
                             "The previous lod_level is {1}; the new "
                             "lod_level is {2}. They are not "
                             "matched".format(self.name, self.lod_level,
@@ -378,18 +395,6 @@ class Variable(object):
                     # get_capacity is implemented
                     pass
 
-        if _in_imperative_mode():
-            # record vars in tracer rather than blocks
-            self._ivar = kwargs.get("ivar", None)
-            if not self._ivar:
-                self._ivar = core.VarBase(stop_gradient)
-            self._ivar.desc = self.desc
-            self._ivar.block = block.desc
-            self._ivar.name = name
-            self._ivar.persistable = persistable
-            if persistable:
-                self.block.vars[name] = self
-        else:
             self.block.vars[name] = self
             self.op = None
             self.stop_gradient = stop_gradient
@@ -462,39 +467,62 @@ class Variable(object):
     def _stop_gradient(self, s):
         if _in_imperative_mode():
             self._ivar.stop_gradient = s
+        else:
             self.stop_gradient = s
 
     @property
     def persistable(self):
+        if _in_imperative_mode():
+            return self._ivar.persistable
+        else:
             return self.desc.persistable()
 
     @persistable.setter
     def persistable(self, p):
+        if _in_imperative_mode():
+            return self._ivar.persistable
+        else:
             self.desc.set_persistable(p)
 
     @property
     def name(self):
+        if _in_imperative_mode():
+            return self._ivar.name
+        else:
             return cpt.to_text(self.desc.name())
 
     @name.setter
     def name(self, new_name):
+        if _in_imperative_mode():
+            self._ivar.name = new_name
+        else:
             self.desc.set_name(new_name)
 
     @property
     def shape(self):
         # convert to tuple, make it as same as numpy API.
+        if _in_imperative_mode():
+            return self._ivar.shape
+        else:
             return tuple(self.desc.shape())
 
     @property
     def dtype(self):
+        if _in_imperative_mode():
+            return self._ivar.dtype
+        else:
             return self.desc.dtype()
 
     @property
     def lod_level(self):
+        # TODO(minqiyang): Support lod_level in imperative mode
         return self.desc.lod_level()
 
     @property
     def type(self):
+        if _in_imperative_mode():
+            return self._ivar.dtype
+        else:
             return self.desc.type()
 
     def _set_error_clip(self, error_clip):
@@ -624,6 +652,32 @@ class Operator(object):
                  inputs=None,
                  outputs=None,
                  attrs=None):
+        if _in_imperative_mode():
+            if type is None:
+                raise ValueError(
+                    "`type` to initilized an Operator can not be None.")
+            self.iop = core.OpBase(type)
+
+            # TODO(minqiyang): remove these lines after we take apart all
+            # backward grads and forward variables
+            self.inputs = defaultdict(list)
+            if inputs is not None:
+                for k, v in six.iteritems(inputs):
+                    if isinstance(v, Variable):
+                        self.inputs[k].append(v._ivar)
+                    elif isinstance(v, list) or isinstance(v, tuple):
+                        self.inputs[k].extend([var._ivar for var in v])
+
+            self.outputs = defaultdict(list)
+            if outputs is not None:
+                for k, v in six.iteritems(outputs):
+                    if isinstance(v, Variable):
+                        self.outputs[k].append(v._ivar)
+                    elif isinstance(v, list) or isinstance(v, tuple):
+                        self.outputs[k].extend([var._ivar for var in v])
+
+            self.attrs = attrs if attrs else {}
+        else:
             self.block = block
             self.desc = desc
             # note: not add self.attrs here:
@@ -636,7 +690,8 @@ class Operator(object):
             op_maker = core.op_proto_and_checker_maker
 
             if op_maker.kOpRoleAttrName() not in op_attrs:
-            op_attrs[op_maker.kOpRoleAttrName()] = self.block.program.op_role
+                op_attrs[op_maker.kOpRoleAttrName(
+                )] = self.block.program.op_role
 
             role_var_name = op_maker.kOpRoleVarAttrName()
             if len(self.block.program.
@@ -699,9 +754,9 @@ class Operator(object):
                     if (m.name not in outputs) and m.dispensable:
                         continue
                     if not ((m.name in outputs) or m.dispensable):
-                    raise ValueError(
-                        ("Incorrect setting for output(s) of "
-                         "operator \"%s\", should set: [%s].") % (type, m.name))
+                        raise ValueError(("Incorrect setting for output(s) of "
+                                          "operator \"%s\", should set: [%s].")
+                                         % (type, m.name))
                 for out_proto in proto.outputs:
                     if out_proto.name not in outputs:
                         continue
@@ -710,8 +765,8 @@ class Operator(object):
                         out_args = [out_args]
                     if not out_proto.duplicable and len(out_args) > 1:
                         raise ValueError(
-                        "Output %s expects only one output, but %d are given." %
-                        (out_proto.name, len(out_args)))
+                            "Output %s expects only one output, but %d are given."
+                            % (out_proto.name, len(out_args)))
                     out_arg_names = []
                     for arg in out_args:
                         out_arg_names.append(cpt.to_text(arg.name))
@@ -725,7 +780,8 @@ class Operator(object):
                     raise TypeError("'attrs' should be a dict.")
                 for attr in proto.attrs:
                     attr_name = attr.name
-                if (attr_name not in op_attrs) or (op_attrs[attr_name] is None):
+                    if (attr_name not in op_attrs) or (
+                            op_attrs[attr_name] is None):
                         continue
                     attr_val = op_attrs[attr_name]
                     self._update_desc_attr(attr_name, attr_val)
@@ -735,26 +791,6 @@ class Operator(object):
                 self.desc.infer_var_type(self.block.desc)
                 self.desc.infer_shape(self.block.desc)
 
-        if _in_imperative_mode():
-            self.iop = core.OpBase()
-            self.iop.desc = self.desc
-
-            self.inputs = defaultdict(list)
-            if inputs is not None:
-                for k, v in six.iteritems(inputs):
-                    if isinstance(v, Variable):
-                        self.inputs[k].append(v._ivar)
-                    elif isinstance(v, list) or isinstance(v, tuple):
-                        self.inputs[k].extend([var._ivar for var in v])
-
-            self.outputs = defaultdict(list)
-            if outputs is not None:
-                for k, v in six.iteritems(outputs):
-                    if isinstance(v, Variable):
-                        self.outputs[k].append(v._ivar)
-                    elif isinstance(v, list) or isinstance(v, tuple):
-                        self.outputs[k].extend([var._ivar for var in v])
-
     def _has_kernel(self, op_type):
         return op_type not in self.OP_WITHOUT_KERNEL_SET
 
@@ -1318,16 +1354,15 @@ class Block(object):
         Returns:
             Operator: the append Operator.
         """
-        op_desc = self.desc.append_op()
+        if _in_imperative_mode():
             op = Operator(
                 block=self,
-            desc=op_desc,
+                desc=None,
                 type=kwargs.get("type", None),
                 inputs=kwargs.get("inputs", None),
                 outputs=kwargs.get("outputs", None),
                 attrs=kwargs.get("attrs", None))
 
-        if _in_imperative_mode():
             # record ops in tracer rather than blocks
             #
             # TODO(minqiyang): add op stop_gradient support in static mode too.
@@ -1335,6 +1370,15 @@ class Block(object):
             _imperative_tracer().trace_op(op,
                                           kwargs.get("stop_gradient", False))
         else:
+            op_desc = self.desc.append_op()
+            op = Operator(
+                block=self,
+                desc=op_desc,
+                type=kwargs.get("type", None),
+                inputs=kwargs.get("inputs", None),
+                outputs=kwargs.get("outputs", None),
+                attrs=kwargs.get("attrs", None))
+
             self.ops.append(op)
 
         return op
@@ -1383,19 +1427,27 @@ class Block(object):
         return self.ops[start:end]
 
     def _prepend_op(self, *args, **kwargs):
-        op_desc = self.desc._prepend_op()
+        if _in_imperative_mode():
             op = Operator(
                 self,
-            op_desc,
+                None,
                 type=kwargs.get("type", None),
                 inputs=kwargs.get("inputs", None),
                 outputs=kwargs.get("outputs", None),
                 attrs=kwargs.get("attrs", None))
-        if _in_imperative_mode():
             _imperative_tracer().trace_op(op,
                                           kwargs.get("stop_gradient", False))
         else:
+            op_desc = self.desc._prepend_op()
+            op = Operator(
+                self,
+                op_desc,
+                type=kwargs.get("type", None),
+                inputs=kwargs.get("inputs", None),
+                outputs=kwargs.get("outputs", None),
+                attrs=kwargs.get("attrs", None))
             self.ops.insert(0, op)
+
         return op
 
     def _sync_with_cpp(self):

python/paddle/fluid/imperative/layers.py

@@ -258,7 +258,7 @@ class PyLayer(core.PyLayer):
             cls.backward_id = core.PyLayer.num_funcs() + 1
             PyLayer.register_func(cls.backward_id, cls._do_backward)
 
-        iop = core.OpBase()
+        iop = core.OpBase(cls.__class__.__name__ + str(cls.forward_id))
         iop.forward_id = cls.forward_id
         iop.backward_id = cls.backward_id
         block.ops.append(iop)

python/paddle/fluid/imperative/tracer.py

@@ -36,14 +36,21 @@ class Tracer(core.Tracer):
         super(Tracer, self).__init__(block)
 
         self._ops = defaultdict()
+        self._vars = defaultdict()
         self._trace_id = 0
 
+    def trace_var(self, name, var):
+        self._vars[name] = var
+
+    def all_parameters(self):
+        return list((item for name, item in six.iteritems(self._vars)
+                     if isinstance(item, framework.Parameter)))
+
     def trace_op(self, op, stop_gradient=False):
         # record op's trace id
         op.iop._trace_id = self._trace_id
 
-        # trace op and save it
-        backward_refs = self.trace(op.iop, op.inputs, op.outputs, op.block.desc,
+        backward_refs = self.trace(op.iop, op.inputs, op.outputs, op.attrs,
                                    framework._current_expected_place(),
                                    stop_gradient)
 

python/paddle/fluid/optimizer.py

@@ -377,17 +377,16 @@ class Optimizer(object):
             and list of (param, grad) Variables pair for optimization.
         """
         self._dtype = loss.dtype
-        program = loss.block.program
         optimize_ops = []
         if framework._in_imperative_mode():
             if parameter_list is not None:
                 parameters = parameter_list
             else:
-                parameters = program.global_block().all_parameters()
+                parameters = framework._imperative_tracer().all_parameters()
 
             params_grads = []
             for param in parameters:
-                if param.stop_gradient or not param.trainable:
+                if not param.trainable:
                     continue
                 # create gradient variable
                 grad_var = Variable(
@@ -396,9 +395,11 @@ class Optimizer(object):
                     stop_gradient=True,
                     ivar=param._ivar._grad_ivar())
                 params_grads.append((param, grad_var))
-            with program_guard(program, startup_program):
+            with program_guard(framework.default_main_program(),
+                               framework.default_startup_program()):
                 optimize_ops = self._create_optimization_pass(params_grads)
         else:
+            program = loss.block.program
             with program_guard(program, startup_program):
                 params_grads = self.backward(loss, startup_program,
                                              parameter_list, no_grad_set)

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

@@ -152,7 +152,7 @@ class SimpleRNNCell(fluid.imperative.Layer):
             type='reduce_sum',
             inputs={'X': softmax_out},
             outputs={'Out': reduce_out},
-            attrs={'dim': None,
+            attrs={'dim': [],
                    'keep_dim': False,
                    'reduce_all': True})
 

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

@@ -277,7 +277,7 @@ class TestImperativeResnet(unittest.TestCase):
 
                 dy_grad_value = {}
                 for param in resnet.parameters():
-                    if not param.stop_gradient:
+                    if param.trainable:
                         np_array = np.array(param._ivar._grad_ivar().value()
                                             .get_tensor())
                         dy_grad_value[param.name + core.grad_var_suffix(
@@ -322,7 +322,7 @@ class TestImperativeResnet(unittest.TestCase):
             for param in resnet.parameters():
                 static_param_name_list.append(param.name)
             for param in resnet.parameters():
-                if not param.stop_gradient:
+                if param.trainable:
                     static_grad_name_list.append(param.name +
                                                  core.grad_var_suffix())