Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
prelu_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_ACTIVATION_PRELU_OP_H_
14 #define DRAGON_OPERATORS_ACTIVATION_PRELU_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 template <class Context>
21 class PReluOp final : public Operator<Context> {
22  public:
23  PReluOp(const OperatorDef& def, Workspace* ws)
24  : Operator<Context>(def, ws),
25  channel_shared_(OpArg<bool>(
26  "channel_shared", false)) {}
28 
29  void RunOnDevice() override;
30  template <typename T> void RunImpl();
31 
32  protected:
34 };
35 
36 template <class Context>
37 class PReluGradientOp final : public Operator<Context> {
38  public:
39  PReluGradientOp(const OperatorDef& def, Workspace* ws)
40  : Operator<Context>(def, ws),
41  channel_shared_(OpArg<bool>(
42  "channel_shared", false)) {}
44 
45  void RunOnDevice() override;
46  template <typename T> void RunImpl();
47 
48  protected:
50 };
51 
52 } // namespace dragon
53 
54 #endif // DRAGON_OPERATORS_ACTIVATION_PRELU_OP_H_
int64_t dim_
Definition: prelu_op.h:49
void RunImpl()
Definition: prelu_op.cc:10
Definition: workspace.h:20
int64_t dim_
Definition: prelu_op.h:33
void RunOnDevice() override
Implement the detailed execution.
Definition: prelu_op.cc:84
int64_t channel_shared_
Definition: prelu_op.h:49
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 channel_shared_
Definition: prelu_op.h:33
Definition: operator.h:149
#define OpArg
Definition: operator.h:235
void RunOnDevice() override
Implement the detailed execution.
Definition: prelu_op.cc:32
USE_OPERATOR_FUNCTIONS
Definition: prelu_op.h:27
PReluOp(const OperatorDef &def, Workspace *ws)
Definition: prelu_op.h:23
Definition: prelu_op.h:21
void RunImpl()
Definition: prelu_op.cc:48
Definition: prelu_op.h:37
int64_t channels_
Definition: prelu_op.h:33
USE_OPERATOR_FUNCTIONS
Definition: prelu_op.h:43
int64_t channels_
Definition: prelu_op.h:49
Definition: common.h:41
PReluGradientOp(const OperatorDef &def, Workspace *ws)
Definition: prelu_op.h:39