add FSetInputVariableAttrs (apache#92)

* add FSetInputVariableAttrs * rename
piiswrong · Jan 7, 2017 · 7b0ac8f · 7b0ac8f
1 parent d6a6e0e
commit 7b0ac8f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/include/nnvm/op_attr_types.h b/include/nnvm/op_attr_types.h
@@ -142,6 +142,18 @@ using FGradient = std::function<std::vector<NodeEntry>(
     const NodePtr& nodeptr,
     const std::vector<NodeEntry>& out_grads)>;
 
+/*!
+ * \brief Set the attributes of input variable.
+ *  Usually used for setting initialization or weight decay.
+ *  \param attrs The attributes of this node.
+ *  \param var the input variable
+ *  \param index index of var in all inputs
+ */
+using FSetInputVarAttrOnCompose = std::function<void(
+    const NodeAttrs& attrs,
+    NodePtr var,
+    const int index)>;
+
 }  // namespace nnvm
 
 #endif  // NNVM_OP_ATTR_TYPES_H_
diff --git a/src/core/symbolic.cc b/src/core/symbolic.cc
@@ -259,6 +259,7 @@ void Symbol::Compose(const array_view<const Symbol*>& args,
                      const std::unordered_map<std::string, const Symbol*>& kwargs,
                      const std::string& name) {
   static auto& flist_inputs = Op::GetAttr<FListInputNames>("FListInputNames");
+  static auto& fset_attrs = Op::GetAttr<FSetInputVarAttrOnCompose>("FSetInputVarAttrOnCompose");
 
   CHECK(!outputs[0].node->is_variable()) << "Variable cannot be composed";
   // parameter check.
@@ -323,6 +324,15 @@ void Symbol::Compose(const array_view<const Symbol*>& args,
       }
     }
     UpdateNodeVersion(n);
+
+    FSetInputVarAttrOnCompose fn = fset_attrs.get(n->op(), nullptr);
+    if (fn != nullptr) {
+      for (size_t i = 0; i < n->inputs.size(); ++i) {
+        if (n->inputs[i].node->is_variable()) {
+          fn(n->attrs, n->inputs[i].node, i);
+        }
+      }
+    }
   } else {
     // general composition
     CHECK_EQ(args.size(), 0U)