Merge pull request #23 from mavenlin/master

support backward of backward

include/mxnet/operator.h

@@ -353,7 +353,7 @@ class OperatorProperty {
     for (size_t i = 0; i < ret_index.size(); ++i) {
       ret[i] = all_data[ret_index[i]];
     }
-    return std::move(ret);
+    return ret;
   }
   /*!
    * \brief create OperatorProperty

include/mxnet/symbolic.h

@@ -85,27 +85,36 @@ class StaticGraph {
     /*! \brief inputs (node_id, index) for of the nodes*/
     std::vector<DataEntry> inputs;
     /*!
-     * \brief If this field is nonnegative, this indicates this
-     *  Node is corresponds to a Backward Operation of Operator.
-     *  backward_source_id will points to the corresponding Forward Node.
+     * \brief source node id; if this field is negative, it means this
+     *  Node is a forward node. If this field is nonnegative, it
+     *  means this Node is the gradient of the source node.
+     */
+    int32_t source_id;
+    /*!
+     * \brief backward; if this field is true, that means this node
+     *  represents the backward function of the op. Else, it
+     *  represents the forward function.  When it represents the
+     *  backward function, itself has not op but shares from the
+     *  source node. It is because the backward function shares the
+     *  states from the forward, and they need to share op.
      *
-     *  For normal node, this field is -1.
-     *  When the node is a Backward node, the op field will be nullptr
+     * Since we support gradient of gradient, a forward node can also
+     * be the gradient of another node. See source id.
      */
-    int32_t backward_source_id;
+    bool backward;
     /*! \brief default constructor */
-    Node() : backward_source_id(-1) {}
+    Node() : source_id(-1), backward(false) {}
     /*! \return whether the node is forward op node */
     inline bool is_forward() const {
-      return op != nullptr;
+      return !backward && !is_variable();
     }
     /*! \return whether the node is backward op node */
     inline bool is_backward() const {
-      return backward_source_id != -1;
+      return backward;
     }
     /*! \return whether the node is variable node */
     inline bool is_variable() const {
-      return op == nullptr && !is_backward();
+      return op == nullptr && source_id == -1;
     }
   };
   /*! \brief all nodes in the graph */

src/symbol/graph_executor.cc

@@ -80,7 +80,7 @@ GraphExecutor::GetResource(uint32_t node_id) const {
     return node.op->ForwardResource();
   } else {
     CHECK(node.is_backward());
-    return graph_.nodes[node.backward_source_id].op->BackwardResource();
+    return graph_.nodes[node.source_id].op->BackwardResource();
   }
 }
 
@@ -90,7 +90,7 @@ inline int GraphExecutor::GetNumOutputs(uint32_t node_id) const {
     return node.op->NumReturns();
   } else if (node.is_backward()) {
     return static_cast<int>(
-        graph_.nodes[node.backward_source_id].op->ListArguments().size());
+        graph_.nodes[node.source_id].op->ListArguments().size());
   } else {
     CHECK(node.is_variable());
     return 1;
@@ -121,11 +121,11 @@ inline std::vector<std::pair<T, T> > GraphExecutor::GetInplaceOption(
       remap[i].first = in_data[rmap_index[i].first];
       remap[i].second = *static_cast<const T*>(rmap_index[i].second);
     }
-    return std::move(remap);
+    return remap;
   } else {
     CHECK(node.is_backward());
     // forward property
-    const OperatorProperty *fwd = graph_.nodes[node.backward_source_id].op.get();
+    const OperatorProperty *fwd = graph_.nodes[node.source_id].op.get();
 
     std::vector<int> out_grad_index(fwd->NumVisibleReturns());
     std::vector<int> in_data_index(fwd->ListArguments().size());
@@ -161,7 +161,7 @@ inline std::vector<std::pair<T, T> > GraphExecutor::GetInplaceOption(
       remap[i].first = *args_array[remap_index[i].first];
       remap[i].second = *static_cast<T*>(remap_index[i].second);
     }
-    return std::move(remap);
+    return remap;
   }
 }
 
@@ -196,7 +196,7 @@ GraphExecutor::GetOpExecEntry(uint32_t nid) {
     op_ctx_ptr->run_ctx = ctx;
     op->Forward(*op_ctx_ptr, in_data, req, out_data);
   };
-  return std::move(exec);
+  return exec;
 }
 
 void GraphExecutor::InitGraph(Symbol symbol, Context ctx, bool need_backward) {
@@ -406,8 +406,8 @@ void GraphExecutor::InitOpNodes() {
     } else {
       CHECK(graph_.nodes[nid].is_backward());
       op_node.op.reset(new BackwardOpWrapper(
-          graph_.nodes[graph_.nodes[nid].backward_source_id].op.get(),
-          op_nodes_[graph_.nodes[nid].backward_source_id].op));
+          graph_.nodes[graph_.nodes[nid].source_id].op.get(),
+          op_nodes_[graph_.nodes[nid].source_id].op));
     }
     bool allow_cache = true;
     for (StaticGraph::DataEntry e : graph_.nodes[nid].inputs) {

src/symbol/static_graph.cc

@@ -18,7 +18,7 @@ std::vector<uint32_t> StaticGraph::TopoSort() const {
       ++out_degree[e.source_id];
     }
     if (n.is_backward()) {
-      ++out_degree[n.backward_source_id];
+      ++out_degree[n.source_id];
     }
   }
   std::vector<uint32_t> ret(nodes.size());
@@ -41,12 +41,12 @@ std::vector<uint32_t> StaticGraph::TopoSort() const {
       }
     }
     if (n.is_backward()) {
-      if (--out_degree[n.backward_source_id] == 0) {
-        queue.push(n.backward_source_id);
+      if (--out_degree[n.source_id] == 0) {
+        queue.push(n.source_id);
       }
     }
   }
-  return std::move(ret);
+  return ret;
 }
 
 bool StaticGraph::InferNodeShapes(const std::vector<uint32_t> &topo_order,
@@ -79,7 +79,7 @@ bool StaticGraph::InferNodeShapes(const std::vector<uint32_t> &topo_order,
       }
     } else if (nodes[nid].is_backward()) {
       // simply use shapes from forward pass to assign backward shape
-      const Node& forward = nodes[node.backward_source_id];
+      const Node& forward = nodes[node.source_id];
       CHECK(forward.is_forward());
       std::vector<TShape>& in_grad_shapes = (*node_out_shapes)[nid];
       CHECK(in_grad_shapes.size() == forward.inputs.size());
@@ -99,7 +99,7 @@ bool StaticGraph::InferNodeShapes(const std::vector<uint32_t> &topo_order,
         }
       }
       // consistent check for input shapes
-      auto& out_data_shapes = (*node_out_shapes)[node.backward_source_id];
+      auto& out_data_shapes = (*node_out_shapes)[node.source_id];
       // use BackwardInputs to select entries corresponding to node.inputs
       auto in_shape = forward.op->BackwardInputs(
           out_data_shapes, in_grad_shapes, out_data_shapes);
@@ -130,7 +130,7 @@ bool StaticGraph::InferShape(std::vector<TShape> *in_shape,
     if (nodes[i].is_forward()) {
       nout = nodes[i].op->NumReturns();
     } else if (nodes[i].is_backward()) {
-      nout = static_cast<int>(nodes[nodes[i].backward_source_id].inputs.size());
+      nout = static_cast<int>(nodes[nodes[i].source_id].inputs.size());
     }
     node_out_shapes[i].resize(nout);
   }
@@ -161,7 +161,7 @@ StaticGraph::Node StaticGraph::CreateSumNode(
   os_size << grad_source.size();
   agg_node.op->Init({{"size", os_size.str()}});
   agg_node.inputs = grad_source;
-  return std::move(agg_node);
+  return agg_node;
 }
 
 void StaticGraph::MakeBackwardPass(std::vector<uint32_t> *head_grad_nodes,
@@ -198,7 +198,6 @@ void StaticGraph::MakeBackwardPass(std::vector<uint32_t> *head_grad_nodes,
     uint32_t nid = *it;
     // skip variables
     if (nodes[nid].is_variable()) continue;
-    CHECK(nodes[nid].is_forward()) << "Do not support Backward of Backward";
     // get out_grad and out_data entry
     std::vector<DataEntry> out_grad, out_data;
     // nvisible is out_grad.size()
@@ -229,7 +228,13 @@ void StaticGraph::MakeBackwardPass(std::vector<uint32_t> *head_grad_nodes,
     // Create a gradient backward node
     Node grad_node;
     // Point to the corresponding source
-    grad_node.backward_source_id = nid;
+    grad_node.source_id = nid;
+    // reverse the source node
+    grad_node.backward = !(nodes[grad_node.source_id].backward);
+    // if grad node is a forward node, needs to have its own OpProperty
+    if (!grad_node.backward) {
+      grad_node.op.reset(nodes[nodes[nid].source_id].op->Copy());
+    }
     // select out the dependent inputs
     grad_node.inputs = nodes[nid].op->BackwardInputs(
         out_grad, nodes[nid].inputs, out_data);

src/symbol/symbol.cc

@@ -1,7 +1,7 @@
 /*!
- *  Copyright (c) 2015 by Contributors
-  *\file symbol.cc
-  *\brief symbol of mxnet
+ * Copyright (c) 2015 by Contributors
+ * \file symbol.cc
+ * \brief symbol of mxnet
  */
 #include <dmlc/logging.h>
 #include <mxnet/symbolic.h>
@@ -12,13 +12,13 @@
 
 namespace mxnet {
 /*!
-  *\brief Node is represents node of an operator in the symbolic graph.
+ * \brief Node is represents node of an operator in the symbolic graph.
  *
-  *It stores connection to the inputs to function represented by OperatorProperty
-  *NOTE on data structure: there are three types of node:
-  *- Normal node: contains all the necessary elements of a graph.
-  *- OperatorProperty: the inputs_ is empty, represents an OperatorProperty that has not been applied.
-  *- Variable: the sym_ is nullptr, represents an named Variable of tensors that can be composed.
+ * It stores connection to the inputs to function represented by OperatorProperty
+ * NOTE on data structure: there are three types of node:
+ * - Normal node: contains all the necessary elements of a graph.
+ * - OperatorProperty: the inputs_ is empty, represents an OperatorProperty that has not been applied.
+ * - Variable: the sym_ is nullptr, represents an named Variable of tensors that can be composed.
  */
 struct Symbol::Node {
   /*! \brief Operator of this node */
@@ -201,7 +201,7 @@ std::vector<std::string> Symbol::ListReturns() const {
       }
     }
   }
-  return std::move(ret);
+  return ret;
 }
 
 Symbol Symbol::operator[] (size_t index) const {
@@ -415,13 +415,13 @@ Symbol Symbol::CreateGroup(const std::vector<Symbol> &symbols) {
   for (const auto &s : symbols) {
     ret.heads_.insert(ret.heads_.end(), s.heads_.begin(), s.heads_.end());
   }
-  return std::move(ret);
+  return ret;
 }
 
 Symbol Symbol::CreateVariable(const std::string &name) {
   Symbol s;
   s.heads_.push_back(DataEntry(std::make_shared<Node>(nullptr, name), 0));
-  return std::move(s);
+  return s;
 }
 
 void Symbol::ToStaticGraph(StaticGraph *out_graph) const {