Caffe2 - C++ API
A deep learning, cross platform ML framework
reduction_front_back_ops.h
1 
17 #ifndef CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_
18 #define CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/logging.h"
22 #include "caffe2/core/operator.h"
23 
24 namespace caffe2 {
25 
26 template <class Context, bool FIRSTDIMS, bool NORMALIZE>
27 class SumReduceDimsOp final : public Operator<Context> {
28  public:
29  SumReduceDimsOp(const OperatorDef& operator_def, Workspace* ws)
30  : Operator<Context>(operator_def, ws),
31  num_reduce_dims_(
32  OperatorBase::GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
33 
34  USE_OPERATOR_CONTEXT_FUNCTIONS;
35 
36  bool RunOnDevice() override {
38  this, Input(0));
39  }
40 
41  template <typename T>
42  bool DoRunWithType() {
43  auto& X = Input(0);
44  auto* Y = Output(0);
45 
46  CAFFE_ENFORCE(
47  num_reduce_dims_ >= 0 && num_reduce_dims_ <= X.dims().size(),
48  "For N-dim input tensor, support num_reduce_dims in range [0, N].");
49 
50  vector<TIndex> output_shape;
51  int start_index = FIRSTDIMS ? num_reduce_dims_ : 0;
52  int end_index =
53  FIRSTDIMS ? X.dims().size() : X.dims().size() - num_reduce_dims_;
54  for (int i = start_index; i < end_index; ++i) {
55  output_shape.push_back(X.dims()[i]);
56  }
57  Y->Resize(output_shape);
58 
59  const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
60  : X.size_to_dim(X.ndim() - num_reduce_dims_);
61  const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
62  : X.size_from_dim(X.ndim() - num_reduce_dims_);
63 
64  if (cols == 0 || rows == 0) {
65  return true;
66  }
67 
68  const T* in_data = X.template data<T>();
69  T* out_data = Y->template mutable_data<T>();
70  Compute(rows, cols, in_data, out_data);
71 
72  return true;
73  }
74 
75  private:
76  template <typename T>
77  void Compute(int rows, int cols, const T* in_data, T* out_data);
78  int num_reduce_dims_;
79 };
80 
81 template <class Context, bool FIRSTDIMS, bool NORMALIZE>
82 class SumReduceDimsGradientOp final : public Operator<Context> {
83  public:
84  SumReduceDimsGradientOp(const OperatorDef& operator_def, Workspace* ws)
85  : Operator<Context>(operator_def, ws),
86  num_reduce_dims_(
87  OperatorBase::GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
88 
89  USE_OPERATOR_CONTEXT_FUNCTIONS;
90 
91  bool RunOnDevice() override {
93  this, Input(0));
94  }
95 
96  template <typename T>
97  bool DoRunWithType() {
98  auto& dY = Input(0);
99  auto& input_1 = Input(1);
100  auto* dX = Output(0);
101 
102  // In previous diff we changed the semantic: Input(1) was changed from
103  // the shape of the input to the data tensor. This made the backward
104  // computation incompatible with old models. To fix this, we check
105  // the dimension and type of Input(1).
106  if (input_1.ndim() == 1 && input_1.template IsType<TIndex>()) {
107  // Input(1) is the shape of the input
108  shape_.CopyFrom(input_1);
109  // Copy first dims
110  vector<TIndex> output_shape(
111  shape_.template data<TIndex>(),
112  shape_.template data<TIndex>() + shape_.size());
113  dX->Resize(output_shape);
114  } else {
115  // Input(1) is data tensor X
116  dX->ResizeLike(input_1);
117  }
118 
119  const int rows = FIRSTDIMS ? dX->size_to_dim(num_reduce_dims_)
120  : dX->size_to_dim(dX->ndim() - num_reduce_dims_);
121  const int cols = FIRSTDIMS
122  ? dX->size_from_dim(num_reduce_dims_)
123  : dX->size_from_dim(dX->ndim() - num_reduce_dims_);
124 
125  const T* dYdata = dY.template data<T>();
126  T* dXdata = dX->template mutable_data<T>();
127  Compute<T>(rows, cols, dYdata, dXdata);
128  return true;
129  }
130 
131  private:
132  template <typename T>
133  void Compute(int rows, int cols, const T* dYdata, T* dXdata);
134  int num_reduce_dims_;
135  // scratch space used for former version of this reducer
136  Tensor<CPUContext> shape_;
137 };
138 
139 template <typename T, class Context, bool FIRSTDIMS>
140 class MaxReduceDimsOp final : public Operator<Context> {
141  public:
142  MaxReduceDimsOp(const OperatorDef& operator_def, Workspace* ws)
143  : Operator<Context>(operator_def, ws),
144  num_reduce_dims_(
145  OperatorBase::GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
146 
147  USE_OPERATOR_CONTEXT_FUNCTIONS;
148 
149  bool RunOnDevice() {
150  auto& X = Input(0);
151  auto* Y = Output(0);
152 
153  CAFFE_ENFORCE(
154  num_reduce_dims_ >= 0 && num_reduce_dims_ <= X.dims().size(),
155  "For N-dim input tensor, support num_reduce_dims in range [0, N].");
156 
157  const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
158  : X.size_to_dim(X.ndim() - num_reduce_dims_);
159  const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
160  : X.size_from_dim(X.ndim() - num_reduce_dims_);
161 
162  vector<TIndex> output_shape;
163  int start_index = FIRSTDIMS ? num_reduce_dims_ : 0;
164  int end_index =
165  FIRSTDIMS ? X.dims().size() : X.dims().size() - num_reduce_dims_;
166 
167  for (int i = start_index; i < end_index; ++i) {
168  output_shape.push_back(X.dims()[i]);
169  }
170  Y->Resize(output_shape);
171 
172  if (cols == 0 || rows == 0) {
173  return true;
174  }
175 
176  const float* data = X.template data<float>();
177  float* out_data = Y->template mutable_data<float>();
178  Compute(rows, cols, data, out_data);
179  return true;
180  }
181 
182  protected:
183  void Compute(int rows, int cols, const float* data, float* out_data);
184 
185  int num_reduce_dims_;
186 };
187 
188 template <typename T, class Context, bool FIRSTDIMS>
189 class MaxReduceDimsGradientOp final : public Operator<Context> {
190  public:
191  MaxReduceDimsGradientOp(const OperatorDef& operator_def, Workspace* ws)
192  : Operator<Context>(operator_def, ws),
193  num_reduce_dims_(
194  OperatorBase::GetSingleArgument<int32_t>("num_reduce_dim", 1)) {}
195 
196  USE_OPERATOR_CONTEXT_FUNCTIONS;
197 
198  bool RunOnDevice() override {
199  auto& dY = Input(0);
200  auto& X = Input(1);
201  auto& Y = Input(2);
202  auto* dX = Output(0);
203 
204  dX->ResizeLike(X);
205  const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
206  : X.size_to_dim(X.ndim() - num_reduce_dims_);
207  const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
208  : X.size_from_dim(X.ndim() - num_reduce_dims_);
209 
210  const float* dYdata = dY.template data<float>();
211  const float* Xdata = X.template data<float>();
212  const float* Ydata = Y.template data<float>();
213 
214  float* dXdata = dX->template mutable_data<float>();
215  Compute(rows, cols, dYdata, Xdata, Ydata, dXdata);
216  return true;
217  }
218 
219  protected:
220  void Compute(
221  int rows,
222  int cols,
223  const float* dYdata,
224  const float* Xdata,
225  const float* Ydata,
226  float* dXdata);
227 
228  int num_reduce_dims_;
229 };
230 
231 } // namespace caffe2
232 
233 #endif // CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.