Caffe2 - C++ API
A deep learning, cross platform ML framework
reduce_ops.cc
1 #include "caffe2/operators/reduce_ops.h"
2 
3 #include "Eigen/Core"
4 
5 #if !EIGEN_VERSION_AT_LEAST(3, 3, 0)
6 #error "Caffe2 requires Eigen to be at least 3.3.0.";
7 #endif
8 
9 #include <unsupported/Eigen/CXX11/Tensor>
10 
11 namespace caffe2 {
12 
13 // For a Tensor X of n dimensions (dims[0], ..., dims[ndim-1]), given index
14 // is converted to corresponding n-dimensional index, e.g. for X.shape = (2,
15 // 3, 4) the linear index 12 maps to 3-dimensional index (1, 0, 0).
16 vector<TIndex> ConvertFromInputIndex(TIndex index, vector<TIndex>& dims) {
17  TIndex ndim = dims.size();
18  vector<TIndex> nd_idx(ndim);
19 
20  for (TIndex i = ndim - 1; i >= 0 && index > 0; i--) {
21  nd_idx[i] = index % dims[i];
22  index /= dims[i];
23  }
24  return nd_idx;
25 }
26 
27 // For given n-dimensional index (nd_idx[0], ..., nd_idx[dims.size()-1]) and
28 // reduction axes, map the n-dimensional index to the corresponding linear
29 // index in the reduced tensor.
30 TIndex ConvertToOutputIndex(
31  const vector<int>& axes,
32  const vector<TIndex>& nd_idx,
33  vector<TIndex>& dims) {
34  TIndex index = 0;
35  TIndex multiplier = 1;
36  for (TIndex i = dims.size() - 1, j = axes.size() - 1; i >= 0; i--) {
37  if (j >= 0 && axes[j] == i) {
38  j--;
39  } else {
40  index += nd_idx[i] * multiplier;
41  multiplier *= dims[i];
42  }
43  }
44  return index;
45 }
46 
47 template <typename T>
48 inline T Add(T x, T y) {
49  return (x + y);
50 }
51 
52 template <typename T, class Context>
53 void ComputeOp(
54  const T* X_data,
55  const TIndex X_size,
56  vector<TIndex>& dims,
57  T* Y_data,
58  vector<int>& axes,
59  int keepdims,
60  T (*binary_op)(T, T)) {
61  for (TIndex x_idx = 0; x_idx < X_size; x_idx++) {
62  vector<TIndex> nd_idx = ConvertFromInputIndex(x_idx, dims);
63  TIndex y_idx = ConvertToOutputIndex(axes, nd_idx, dims);
64  Y_data[y_idx] = binary_op(Y_data[y_idx], X_data[x_idx]);
65  }
66 }
67 
68 namespace {
69 
70 template <typename U, int DIMS>
71 using ReductionTensor = Eigen::Tensor<U, DIMS, Eigen::RowMajor>;
72 
73 template <int DIMS>
74 using DSizesType = Eigen::DSizes<Eigen::DenseIndex, DIMS>;
75 
76 template <int DIMS>
77 DSizesType<DIMS> calcDSize(vector<TIndex>& dims) {
78  Eigen::DSizes<Eigen::DenseIndex, DIMS> dsizes_out;
79  size_t i = 0;
80  for (i = 0; i < DIMS; ++i) {
81  if (i < dims.size()) {
82  dsizes_out[i] = dims[i];
83  } else {
84  dsizes_out[i] = 1;
85  }
86  }
87  return dsizes_out;
88 }
89 
90 } // namespace
91 
92 template <typename T, class Context>
93 bool ReduceSumOp<T, Context>::Compute(
94  const T* X_data,
95  const TIndex X_size,
96  vector<TIndex>& dims,
97  T* Y_data,
98  const TIndex Y_size,
99  vector<int>& axes,
100  vector<TIndex>& Y_dims,
101  int keepdims) {
102  switch (dims.size()) {
103  case 1: {
104  std::array<int, 1> reduce_dims{{0}};
105  Eigen::DSizes<Eigen::DenseIndex, 1> dsizes_X = calcDSize<1>(dims);
106  Eigen::DSizes<Eigen::DenseIndex, 1> dsizes_Y = calcDSize<1>(Y_dims);
107  auto X_ten =
108  Eigen::TensorMap<ReductionTensor<const T, 1>>(X_data, dsizes_X);
109  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 1>>(Y_data, dsizes_Y);
110  Y_ten = X_ten.sum(reduce_dims);
111  } break;
112  case 2: {
113  Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_X = calcDSize<2>(dims);
114  Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_Y = calcDSize<2>(Y_dims);
115  auto X_ten =
116  Eigen::TensorMap<ReductionTensor<const T, 2>>(X_data, dsizes_X);
117  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 2>>(Y_data, dsizes_Y);
118  switch (axes.size()) {
119  case 1: {
120  std::array<int, 1> reduce_dims;
121  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
122  Y_ten = X_ten.sum(reduce_dims);
123  } break;
124  case 2: {
125  std::array<int, 2> reduce_dims;
126  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
127  Y_ten = X_ten.sum(reduce_dims);
128  } break;
129  }
130  } break;
131  case 3: {
132  Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_X = calcDSize<3>(dims);
133  Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_Y = calcDSize<3>(Y_dims);
134  auto X_ten =
135  Eigen::TensorMap<ReductionTensor<const T, 3>>(X_data, dsizes_X);
136  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 3>>(Y_data, dsizes_Y);
137  switch (axes.size()) {
138  case 1: {
139  std::array<int, 1> reduce_dims;
140  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
141  Y_ten = X_ten.sum(reduce_dims);
142  } break;
143  case 2: {
144  std::array<int, 2> reduce_dims;
145  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
146  Y_ten = X_ten.sum(reduce_dims);
147  } break;
148  case 3: {
149  std::array<int, 3> reduce_dims;
150  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
151  Y_ten = X_ten.sum(reduce_dims);
152  }
153  }
154  } break;
155  default: {
156  math::Set<T, Context>(Y_size, 0.f, Y_data, &context_);
157  ComputeOp<T, Context>(X_data, X_size, dims, Y_data, axes, keepdims, Add);
158  }
159  }
160  return true;
161 }
162 
163 template <typename T, class Context>
164 bool ReduceMeanOp<T, Context>::Compute(
165  const T* X_data,
166  const TIndex X_size,
167  vector<TIndex>& dims,
168  T* Y_data,
169  const TIndex Y_size,
170  vector<int>& axes,
171  vector<TIndex>& Y_dims,
172  int keepdims) {
173  switch (dims.size()) {
174  case 1: {
175  std::array<int, 1> reduce_dims{{0}};
176  Eigen::DSizes<Eigen::DenseIndex, 1> dsizes_X = calcDSize<1>(dims);
177  Eigen::DSizes<Eigen::DenseIndex, 1> dsizes_Y = calcDSize<1>(Y_dims);
178  auto X_ten =
179  Eigen::TensorMap<ReductionTensor<const T, 1>>(X_data, dsizes_X);
180  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 1>>(Y_data, dsizes_Y);
181  Y_ten = X_ten.mean(reduce_dims);
182  } break;
183  case 2: {
184  Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_X = calcDSize<2>(dims);
185  Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_Y = calcDSize<2>(Y_dims);
186  auto X_ten =
187  Eigen::TensorMap<ReductionTensor<const T, 2>>(X_data, dsizes_X);
188  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 2>>(Y_data, dsizes_Y);
189  switch (axes.size()) {
190  case 1: {
191  std::array<int, 1> reduce_dims;
192  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
193  Y_ten = X_ten.mean(reduce_dims);
194  } break;
195  case 2: {
196  std::array<int, 2> reduce_dims;
197  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
198  Y_ten = X_ten.mean(reduce_dims);
199  } break;
200  }
201  } break;
202  case 3: {
203  Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_X = calcDSize<3>(dims);
204  Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_Y = calcDSize<3>(Y_dims);
205  auto X_ten =
206  Eigen::TensorMap<ReductionTensor<const T, 3>>(X_data, dsizes_X);
207  auto Y_ten = Eigen::TensorMap<ReductionTensor<T, 3>>(Y_data, dsizes_Y);
208  switch (axes.size()) {
209  case 1: {
210  std::array<int, 1> reduce_dims;
211  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
212  Y_ten = X_ten.mean(reduce_dims);
213  } break;
214  case 2: {
215  std::array<int, 2> reduce_dims;
216  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
217  Y_ten = X_ten.mean(reduce_dims);
218  } break;
219  case 3: {
220  std::array<int, 3> reduce_dims;
221  std::copy(axes.begin(), axes.end(), reduce_dims.begin());
222  Y_ten = X_ten.mean(reduce_dims);
223  }
224  }
225  } break;
226  default: {
227  math::Set<T, Context>(Y_size, 0.f, Y_data, &context_);
228  ComputeOp<T, Context>(X_data, X_size, dims, Y_data, axes, keepdims, Add);
229  math::Scale(
230  Y_size,
231  static_cast<float>(Y_size) / X_size,
232  Y_data,
233  Y_data,
234  &context_);
235  } break;
236  }
237 
238  return true;
239 }
240 
241 REGISTER_CPU_OPERATOR(ReduceSum, ReduceSumOp<float, CPUContext>);
242 
243 OPERATOR_SCHEMA(ReduceSum)
244  .NumInputs(1)
245  .NumOutputs(1)
246  .SetDoc(R"DOC(
247  Computes the sum of the input tensor's element along the provided axes.
248  The resulted tensor has the same rank as the input if keepdims equal 1.
249  If keepdims equal 0, then the resulted tensor have the reduced dimension pruned.
250 )DOC")
251  .Arg("axes", "A list of integers, along which to reduce.")
252  .Arg(
253  "keepdims",
254  "Keep the reduced dimension(s) or not, default 1 keeps the reduced dimension(s).")
255  .Input(0, "data", "An input tensor.")
256  .Output(0, "reduced", "Reduced output tensor.");
257 
258 // TODO: Write gradient for this when needed
259 GRADIENT_NOT_IMPLEMENTED_YET(ReduceSum);
260 
261 REGISTER_CPU_OPERATOR(ReduceMean, ReduceMeanOp<float, CPUContext>);
262 
263 OPERATOR_SCHEMA(ReduceMean)
264  .NumInputs(1)
265  .NumOutputs(1)
266  .SetDoc(R"DOC(
267  Computes the mean of the input tensor's element along the provided axes.
268  The resulted tensor has the same rank as the input if keepdims equal 1.
269  If keepdims equal 0, then the resulted tensor have the reduced dimension pruned.
270  )DOC")
271  .Arg("axes", "A list of integers, along which to reduce.")
272  .Arg(
273  "keepdims",
274  "Keep the reduced dimension(s) or not, default 1 keeps the reduced dimension(s).")
275  .Input(0, "data", "An input tensor.")
276  .Output(0, "reduced", "Reduced output tensor.");
277 
278 // TODO: Write gradient for this when needed
279 GRADIENT_NOT_IMPLEMENTED_YET(ReduceMean);
280 
281 } // namespace caffe2
caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: convert_encoded_to_raw_leveldb.cc:47
Generated on Thu Apr 19 2018 13:03:56 for Caffe2 - C++ API by   doxygen 1.8.11