Commit 02396a7f by Eric Junyuan Xie Committed by Tianqi Chen

add FSetInputVariableAttrs (#92)

* add FSetInputVariableAttrs

* rename
parent 84b29b7b
......@@ -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_
......@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment