Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
lrn_op.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_OPERATORS_VISION_LRN_OP_H_
14 #define DRAGON_OPERATORS_VISION_LRN_OP_H_
15 
16 #include "core/operator.h"
17 
18 namespace dragon {
19 
20 typedef enum {
23 } LRNMode;
24 
25 template <class Context>
26 class LRNOp : public Operator<Context> {
27  public:
28  LRNOp(const OperatorDef& def, Workspace* ws)
29  : Operator<Context>(def, ws),
30  local_size_(OpArg<int64_t>(
31  "local_size", 5)),
32  alpha_(OpArg<float>(
33  "alpha", 0.0001f)),
34  beta_(OpArg<float>(
35  "beta", 0.75f)),
36  k_(OpArg<float>(
37  "k", 2.f)),
38  mode_(OpArg<string>(
39  "mode", "ACROSS_CHANNELS")) {}
41 
42  void RunOnDevice() override;
43  template <typename T> void AcrossRunImpl();
44  template <typename T> void SplitRunImpl();
45  template <typename T> void SquareRunImpl();
46  template <typename T> void PoolRunImpl();
47  template <typename T> void PowRunImpl();
48  template <typename T> void ProdRunImpl();
49 
50  protected:
51  string mode_;
53  float alpha_, beta_, k_;
54  unique_ptr<OperatorBase> sqr_op_, pool_op_;
55  unique_ptr<OperatorBase> pow_op_, prod_op_;
58 };
59 
60 template <class Context>
61 class LRNGradientOp : public Operator<Context> {
62  public:
63  LRNGradientOp(const OperatorDef& def, Workspace* ws)
64  : Operator<Context>(def, ws),
65  local_size_(OpArg<int64_t>(
66  "local_size", 5)),
67  alpha_(OpArg<float>(
68  "alpha", 0.0001f)),
69  beta_(OpArg<float>(
70  "beta", 0.75f)),
71  k_(OpArg<float>(
72  "k", 2.f)),
73  mode_(OpArg<string>(
74  "mode", "ACROSS_CHANNELS")) {}
76 
77  void RunOnDevice() override;
78  template <typename T> void AcrossRunImpl();
79  template <typename T> void SplitRunImpl();
80  template <typename T> void SquareRunImpl();
81  template <typename T> void PoolRunImpl();
82  template <typename T> void PowRunImpl();
83  template <typename T> void ProdRunImpl();
84 
85  protected:
86  string mode_;
88  float alpha_, beta_, k_;
89  unique_ptr<OperatorBase> sqr_op_, pool_op_;
90  unique_ptr<OperatorBase> pow_op_, prod_op_;
93 };
94 
95 #ifdef WITH_CUDNN
96 
97 template <class Context>
98 class CuDNNLRNOp final : public LRNOp<Context> {
99  public:
100  CuDNNLRNOp(const OperatorDef& def, Workspace* ws)
101  : LRNOp<Context>(def, ws) {
102  CUDNN_CHECK(cudnnCreateTensorDescriptor(&input_desc_));
103  CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc_));
104  CUDNN_CHECK(cudnnCreateLRNDescriptor(&lrn_desc_));
105  CUDNN_CHECK(cudnnSetLRNDescriptor(
106  lrn_desc_,
107  this->local_size_,
108  this->alpha_,
109  this->beta_,
110  this->k_
111  ));
112  }
114 
116  CUDNN_CHECK(cudnnDestroyTensorDescriptor(input_desc_));
117  CUDNN_CHECK(cudnnDestroyTensorDescriptor(output_desc_));
118  CUDNN_CHECK(cudnnDestroyLRNDescriptor(lrn_desc_));
119  }
120 
121  void RunOnDevice() override;
122  template <typename T> void RunImpl();
123 
124  protected:
125  cudnnTensorDescriptor_t input_desc_, output_desc_;
126  cudnnLRNDescriptor_t lrn_desc_;
127 };
128 
129 template <class Context>
130 class CuDNNLRNGradientOp final : public LRNGradientOp<Context > {
131  public:
132  CuDNNLRNGradientOp(const OperatorDef& def, Workspace* ws)
133  : LRNGradientOp<Context>(def, ws) {
134  CUDNN_CHECK(cudnnCreateTensorDescriptor(&input_desc_));
135  CUDNN_CHECK(cudnnCreateTensorDescriptor(&output_desc_));
136  CUDNN_CHECK(cudnnCreateLRNDescriptor(&lrn_desc_));
137  CUDNN_CHECK(cudnnSetLRNDescriptor(
138  lrn_desc_,
139  this->local_size_,
140  this->alpha_,
141  this->beta_,
142  this->k_
143  ));
144  }
146 
148  CUDNN_CHECK(cudnnDestroyTensorDescriptor(input_desc_));
149  CUDNN_CHECK(cudnnDestroyTensorDescriptor(output_desc_));
150  CUDNN_CHECK(cudnnDestroyLRNDescriptor(lrn_desc_));
151  }
152 
153  void RunOnDevice() override;
154  template <typename T> void RunImpl();
155 
156  protected:
157  cudnnTensorDescriptor_t input_desc_, output_desc_;
158  cudnnLRNDescriptor_t lrn_desc_;
159 };
160 
161 #endif // WITH CUDNN
162 
163 } // namespace dragon
164 
165 #endif // DRAGON_OPERATORS_VISION_LRN_OP_H_
USE_OPERATOR_FUNCTIONS
Definition: lrn_op.h:40
cudnnTensorDescriptor_t input_desc_
Definition: lrn_op.h:125
float beta_
Definition: lrn_op.h:53
Tensor * pool_out_
Definition: lrn_op.h:57
Definition: workspace.h:20
float beta_
Definition: lrn_op.h:88
Definition: lrn_op.h:61
void ProdRunImpl()
Definition: lrn_op.cc:95
LRNOp(const OperatorDef &def, Workspace *ws)
Definition: lrn_op.h:28
#define CUDNN_CHECK(condition)
Definition: cudnn_device.h:34
cudnnTensorDescriptor_t output_desc_
Definition: lrn_op.h:125
~CuDNNLRNGradientOp()
Definition: lrn_op.h:147
void RunImpl()
Definition: cudnn_lrn_op.cc:8
Tensor * prod_in_
Definition: lrn_op.h:56
cudnnTensorDescriptor_t input_desc_
Definition: lrn_op.h:157
void RunOnDevice() override
Implement the detailed execution.
Definition: cudnn_lrn_op.cc:31
void AcrossRunImpl()
Definition: lrn_op.cc:11
const OperatorDef & def() const
Return the stored def.
Definition: operator.h:114
void PoolRunImpl()
Definition: lrn_op.cc:185
Workspace * ws() const
Return the parent workspace.
Definition: operator.h:87
LRNGradientOp(const OperatorDef &def, Workspace *ws)
Definition: lrn_op.h:63
unique_ptr< OperatorBase > sqr_op_
Definition: lrn_op.h:54
Tensor * sqr_in_
Definition: lrn_op.h:91
string mode_
Definition: lrn_op.h:86
void SplitRunImpl()
Definition: lrn_op.cc:17
CuDNNLRNOp(const OperatorDef &def, Workspace *ws)
Definition: lrn_op.h:100
unique_ptr< OperatorBase > pool_op_
Definition: lrn_op.h:89
Definition: lrn_op.h:130
cudnnTensorDescriptor_t output_desc_
Definition: lrn_op.h:157
unique_ptr< OperatorBase > pool_op_
Definition: lrn_op.h:54
void RunOnDevice() override
Implement the detailed execution.
Definition: lrn_op.cc:246
unique_ptr< OperatorBase > sqr_op_
Definition: lrn_op.h:89
int local_size_
Definition: lrn_op.h:87
void RunOnDevice() override
Implement the detailed execution.
Definition: cudnn_lrn_op.cc:72
USE_OPERATOR_FUNCTIONS
Definition: lrn_op.h:75
Definition: tensor.h:21
void SquareRunImpl()
Definition: lrn_op.cc:29
Tensor * pow_out_
Definition: lrn_op.h:57
void ProdRunImpl()
Definition: lrn_op.cc:139
void RunImpl()
Definition: cudnn_lrn_op.cc:45
Definition: operator.h:149
#define OpArg
Definition: operator.h:235
Definition: lrn_op.h:22
Tensor * pow_out_
Definition: lrn_op.h:92
cudnnLRNDescriptor_t lrn_desc_
Definition: lrn_op.h:126
unique_ptr< OperatorBase > prod_op_
Definition: lrn_op.h:55
int local_size_
Definition: lrn_op.h:52
unique_ptr< OperatorBase > pow_op_
Definition: lrn_op.h:55
CuDNNLRNGradientOp(const OperatorDef &def, Workspace *ws)
Definition: lrn_op.h:132
float k_
Definition: lrn_op.h:53
Definition: lrn_op.h:98
string mode_
Definition: lrn_op.h:51
Tensor * sqr_out_
Definition: lrn_op.h:92
Tensor * pool_out_
Definition: lrn_op.h:92
Tensor * scale
Definition: lrn_op.h:56
Definition: lrn_op.h:21
cudnnLRNDescriptor_t lrn_desc_
Definition: lrn_op.h:158
Definition: lrn_op.h:26
unique_ptr< OperatorBase > pow_op_
Definition: lrn_op.h:90
Tensor * prod_in_
Definition: lrn_op.h:91
Tensor * scale_
Definition: lrn_op.h:91
float alpha_
Definition: lrn_op.h:88
void SplitRunImpl()
Definition: lrn_op.cc:231
float k_
Definition: lrn_op.h:88
void AcrossRunImpl()
Definition: lrn_op.cc:133
Tensor * sqr_in_
Definition: lrn_op.h:56
Tensor * sqr_out_
Definition: lrn_op.h:57
USE_OPERATOR_FUNCTIONS
Definition: lrn_op.h:113
float alpha_
Definition: lrn_op.h:53
LRNMode
Definition: lrn_op.h:20
void PoolRunImpl()
Definition: lrn_op.cc:49
unique_ptr< OperatorBase > prod_op_
Definition: lrn_op.h:90
Definition: common.h:41
void RunOnDevice() override
Implement the detailed execution.
Definition: lrn_op.cc:114
~CuDNNLRNOp()
Definition: lrn_op.h:115
USE_OPERATOR_FUNCTIONS
Definition: lrn_op.h:145
void PowRunImpl()
Definition: lrn_op.cc:162
void PowRunImpl()
Definition: lrn_op.cc:73
void SquareRunImpl()
Definition: lrn_op.cc:210