Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
softmax_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_ACTIVATION_SOFTMAX_OP_H_
14 #define DRAGON_OPERATORS_ACTIVATION_SOFTMAX_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 template <class Context>
21 class SoftmaxOp final : public Operator<Context> {
22  public:
23  SoftmaxOp(const OperatorDef& def, Workspace* ws)
24  : Operator<Context>(def, ws),
25  axis_(OpArg<int64_t>("axis", 1)) {}
27 
28  void RunOnDevice() override;
29  template <typename T> void RunImpl();
30 
31  protected:
32  int64_t axis_, axis_dim_;
34 };
35 
36 template <class Context>
37 class SoftmaxGradientOp final : public Operator<Context> {
38  public:
39  SoftmaxGradientOp(const OperatorDef& def, Workspace* ws)
40  : Operator<Context>(def, ws),
41  axis_(OpArg<int64_t>("axis", 1)) {}
43 
44  void RunOnDevice() override;
45  template <typename T> void RunImpl();
46 
47  protected:
48  int64_t axis_, axis_dim_;
50 };
51 
52 #ifdef WITH_CUDNN
53 
54 template <class Context>
55 class CuDNNSoftmaxOp final : public Operator<Context> {
56  public:
57  CuDNNSoftmaxOp(const OperatorDef& def, Workspace* ws)
58  : Operator<Context>(def, ws),
59  axis_(OpArg<int64_t>("axis", 1)) {
61  }
63 
66  }
67 
68  void RunOnDevice() override;
69  template <typename T> void RunImpl();
70 
71  protected:
73  cudnnTensorDescriptor_t input_desc_;
74 };
75 
76 template <class Context>
77 class CuDNNSoftmaxGradientOp final : public Operator<Context> {
78  public:
79  CuDNNSoftmaxGradientOp(const OperatorDef& def, Workspace* ws)
80  : Operator<Context>(def, ws),
81  axis_(OpArg<int64_t>("axis", 1)) {
83  }
85 
88  }
89 
90  void RunOnDevice() override;
91  template <typename T> void RunImpl();
92 
93  protected:
95  cudnnTensorDescriptor_t input_desc_;
96 };
97 
98 #endif // WITH_CUDNN
99 
100 } // namespace dragon
101 
102 #endif // DRAGON_OPERATORS_ACTIVATION_SOFTMAX_OP_H_
int64_t axis_dim_
Definition: softmax_op.h:48
int64_t outer_dim_
Definition: softmax_op.h:33
int64_t axis_
Definition: softmax_op.h:48
void RunOnDevice() override
Implement the detailed execution.
Definition: cudnn_softmax_op.cc:41
Definition: workspace.h:20
cudnnTensorDescriptor_t input_desc_
Definition: softmax_op.h:73
const OperatorDef & def() const
Return the stored def.
Definition: operator.h:114
Workspace * ws() const
Return the parent workspace.
Definition: operator.h:87
USE_OPERATOR_FUNCTIONS
Definition: softmax_op.h:62
void RunImpl()
Definition: cudnn_softmax_op.cc:53
void RunOnDevice() override
Implement the detailed execution.
Definition: cudnn_softmax_op.cc:80
int64_t inner_dim_
Definition: softmax_op.h:94
CuDNNSoftmaxOp(const OperatorDef &def, Workspace *ws)
Definition: softmax_op.h:57
int64_t outer_dim_
Definition: softmax_op.h:72
cudnnTensorDescriptor_t input_desc_
Definition: softmax_op.h:95
Definition: softmax_op.h:55
void CuDNNCreateTensorDesc(cudnnTensorDescriptor_t *desc)
Definition: cudnn_device.h:67
CuDNNSoftmaxGradientOp(const OperatorDef &def, Workspace *ws)
Definition: softmax_op.h:79
SoftmaxOp(const OperatorDef &def, Workspace *ws)
Definition: softmax_op.h:23
void RunImpl()
Definition: cudnn_softmax_op.cc:16
Definition: operator.h:149
#define OpArg
Definition: operator.h:235
void RunImpl()
Definition: softmax_op.cc:51
int64_t inner_dim_
Definition: softmax_op.h:33
void RunImpl()
Definition: softmax_op.cc:16
Definition: softmax_op.h:21
Definition: softmax_op.h:37
Definition: softmax_op.h:77
void RunOnDevice() override
Implement the detailed execution.
Definition: softmax_op.cc:76
USE_OPERATOR_FUNCTIONS
Definition: softmax_op.h:26
void RunOnDevice() override
Implement the detailed execution.
Definition: softmax_op.cc:38
USE_OPERATOR_FUNCTIONS
Definition: softmax_op.h:42
void CuDNNDestroyTensorDesc(cudnnTensorDescriptor_t *desc)
Definition: cudnn_device.h:72
int64_t inner_dim_
Definition: softmax_op.h:72
~CuDNNSoftmaxGradientOp()
Definition: softmax_op.h:86
int64_t outer_dim_
Definition: softmax_op.h:49
int64_t axis_
Definition: softmax_op.h:94
int64_t axis_
Definition: softmax_op.h:32
~CuDNNSoftmaxOp()
Definition: softmax_op.h:64
SoftmaxGradientOp(const OperatorDef &def, Workspace *ws)
Definition: softmax_op.h:39
int64_t axis_
Definition: softmax_op.h:72
Definition: common.h:41
USE_OPERATOR_FUNCTIONS
Definition: softmax_op.h:84
int64_t inner_dim_
Definition: softmax_op.h:49
int64_t axis_dim_
Definition: softmax_op.h:32
int64_t outer_dim_
Definition: softmax_op.h:94