marian-nmt/marian
1
1
Fork 0
mirror of https://github.com/marian-nmt/marian.git synced 2024-09-11 06:15:56 +03:00
Code Issues Packages Projects Releases Wiki Activity

add ArgmaxNodeOp

Browse Source
This commit is contained in:
Hieu Hoang 2016-09-16 19:49:59 +02:00
parent c824de0e82
commit 948789787d
1 changed files with 29 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
src/node_operators.h
View File

@ -178,6 +178,35 @@ struct SoftmaxNodeOp : public UnaryNodeOp {
};
struct ArgmaxNodeOp : public UnaryNodeOp {
template <typename ...Args>
ArgmaxNodeOp(ChainPtr a, Args ...args)
: UnaryNodeOp(a, keywords::shape=newShape(a), args...) { }
void forward() {
// B = softmax(A).
Argmax(&val_, &a_->val());
}
void backward() {
}
Shape newShape(ChainPtr a) {
Shape shape = a->shape();
shape[1] = 1;
return shape;
}
virtual std::string graphviz() {
std::stringstream ss;
ss << "\"" << this << "\" [shape=\"box\", label=\"argmax\", style=\"filled\", fillcolor=\"yellow\"]" << std::endl;
ss << "\"" << &*a_ << "\" -> \"" << this << "\"" << std::endl << std::endl;
return ss.str();
};
};
struct LogNodeOp : public UnaryNodeOp {
template <typename ...Args>
LogNodeOp(Args ...args)