Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
one_hot_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_ARRAY_ONE_HOT_OP_H_
14 #define DRAGON_OPERATORS_ARRAY_ONE_HOT_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 template <class Context>
21 class OneHotOp final : public Operator<Context> {
22  public:
23  OneHotOp(const OperatorDef& def, Workspace* ws)
24  : Operator<Context>(def, ws),
25  depth_(OpArg<int64_t>("depth", -1)),
26  on_value_(OpArg<int64_t>("on_value", 1)),
27  off_value_(OpArg<int64_t>("off_value", 0)) {}
29 
30  void RunOnDevice() override;
31  template <typename T> void RunImpl();
32 
33  protected:
35 };
36 
37 } // namespace dragon
38 
39 #endif // DRAGON_OPERATORS_ARRAY_ONE_HOT_OP_H_
int64_t depth_
Definition: one_hot_op.h:34
Definition: workspace.h:20
OneHotOp(const OperatorDef &def, Workspace *ws)
Definition: one_hot_op.h:23
const OperatorDef & def() const
Return the stored def.
Definition: operator.h:114
Workspace * ws() const
Return the parent workspace.
Definition: operator.h:87
int64_t off_value_
Definition: one_hot_op.h:34
Definition: operator.h:149
#define OpArg
Definition: operator.h:235
Definition: one_hot_op.h:21
USE_OPERATOR_FUNCTIONS
Definition: one_hot_op.h:28
void RunImpl()
Definition: one_hot_op.cc:8
int64_t on_value_
Definition: one_hot_op.h:34
void RunOnDevice() override
Implement the detailed execution.
Definition: one_hot_op.cc:26
Definition: common.h:41