1 #include "caffe2/operators/reduce_ops.h" 5 #if !EIGEN_VERSION_AT_LEAST(3, 3, 0) 6 #error "Caffe2 requires Eigen to be at least 3.3.0."; 9 #include <unsupported/Eigen/CXX11/Tensor> 16 vector<TIndex> ConvertFromInputIndex(TIndex index, vector<TIndex>& dims) {
17 TIndex ndim = dims.size();
18 vector<TIndex> nd_idx(ndim);
20 for (TIndex i = ndim - 1; i >= 0 && index > 0; i--) {
21 nd_idx[i] = index % dims[i];
30 TIndex ConvertToOutputIndex(
31 const vector<int>& axes,
32 const vector<TIndex>& nd_idx,
33 vector<TIndex>& dims) {
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) {
40 index += nd_idx[i] * multiplier;
41 multiplier *= dims[i];
48 inline T Add(T x, T y) {
52 template <
typename T,
class Context>
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]);
70 template <
typename U,
int DIMS>
71 using ReductionTensor = Eigen::Tensor<U, DIMS, Eigen::RowMajor>;
74 using DSizesType = Eigen::DSizes<Eigen::DenseIndex, DIMS>;
77 DSizesType<DIMS> calcDSize(vector<TIndex>& dims) {
78 Eigen::DSizes<Eigen::DenseIndex, DIMS> dsizes_out;
80 for (i = 0; i < DIMS; ++i) {
81 if (i < dims.size()) {
82 dsizes_out[i] = dims[i];
92 template <
typename T,
class Context>
93 bool ReduceSumOp<T, Context>::Compute(
100 vector<TIndex>& Y_dims,
102 switch (dims.size()) {
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);
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);
113 Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_X = calcDSize<2>(dims);
114 Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_Y = calcDSize<2>(Y_dims);
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()) {
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);
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);
132 Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_X = calcDSize<3>(dims);
133 Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_Y = calcDSize<3>(Y_dims);
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()) {
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);
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);
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);
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);
163 template <
typename T,
class Context>
164 bool ReduceMeanOp<T, Context>::Compute(
167 vector<TIndex>& dims,
171 vector<TIndex>& Y_dims,
173 switch (dims.size()) {
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);
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);
184 Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_X = calcDSize<2>(dims);
185 Eigen::DSizes<Eigen::DenseIndex, 2> dsizes_Y = calcDSize<2>(Y_dims);
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()) {
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);
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);
203 Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_X = calcDSize<3>(dims);
204 Eigen::DSizes<Eigen::DenseIndex, 3> dsizes_Y = calcDSize<3>(Y_dims);
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()) {
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);
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);
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);
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);
231 static_cast<float>(Y_size) / X_size,
241 REGISTER_CPU_OPERATOR(ReduceSum, ReduceSumOp<float, CPUContext>);
243 OPERATOR_SCHEMA(ReduceSum)
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. 251 .Arg("axes",
"A list of integers, along which to reduce.")
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.");
259 GRADIENT_NOT_IMPLEMENTED_YET(ReduceSum);
261 REGISTER_CPU_OPERATOR(ReduceMean, ReduceMeanOp<float, CPUContext>);
263 OPERATOR_SCHEMA(ReduceMean)
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. 271 .Arg("axes",
"A list of integers, along which to reduce.")
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.");
279 GRADIENT_NOT_IMPLEMENTED_YET(ReduceMean);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...