reduce_op.h 8.9 KB
Newer Older
G
guosheng 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License. */

#pragma once

D
Dong Zhihong 已提交
17
#include "glog/logging.h"
G
guosheng 已提交
18 19 20 21 22 23 24 25 26 27 28
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using DDim = framework::DDim;
template <typename T, size_t D, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenTensor = framework::EigenTensor<T, D, MajorType, IndexType>;
D
Dong Zhihong 已提交
29 30 31
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenScalar = framework::EigenScalar<T, MajorType, IndexType>;
32 33 34
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;
D
Dong Zhihong 已提交
35

G
guosheng 已提交
36
struct SumFunctor {
Q
QI JUN 已提交
37 38
  template <typename DeviceContext, typename X, typename Y, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, const Dim& dim) {
G
guosheng 已提交
39
    y.device(place) = x.sum(dim);
G
guosheng 已提交
40 41 42 43
  }
};

struct SumGradFunctor {
Q
QI JUN 已提交
44 45 46
  template <typename DeviceContext, typename X, typename Y, typename DX,
            typename DY, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, DX& dx, DY& dy,
G
guosheng 已提交
47 48
                  const Dim& dim, int size) {
    dx.device(place) = dy.broadcast(dim);
G
guosheng 已提交
49 50 51 52
  }
};

struct MeanFunctor {
Q
QI JUN 已提交
53 54
  template <typename DeviceContext, typename X, typename Y, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, const Dim& dim) {
G
guosheng 已提交
55
    y.device(place) = x.mean(dim);
G
guosheng 已提交
56 57 58 59
  }
};

struct MeanGradFunctor {
Q
QI JUN 已提交
60 61 62
  template <typename DeviceContext, typename X, typename Y, typename DX,
            typename DY, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, DX& dx, DY& dy,
G
guosheng 已提交
63 64
                  const Dim& dim, int size) {
    dx.device(place) = dy.broadcast(dim) / dx.constant(size);
G
guosheng 已提交
65 66 67 68
  }
};

struct MaxFunctor {
Q
QI JUN 已提交
69 70
  template <typename DeviceContext, typename X, typename Y, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, const Dim& dim) {
G
guosheng 已提交
71
    y.device(place) = x.maximum(dim);
G
guosheng 已提交
72 73 74 75
  }
};

struct MinFunctor {
Q
QI JUN 已提交
76 77
  template <typename DeviceContext, typename X, typename Y, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, const Dim& dim) {
G
guosheng 已提交
78
    y.device(place) = x.minimum(dim);
G
guosheng 已提交
79 80 81 82
  }
};

struct MaxOrMinGradFunctor {
Q
QI JUN 已提交
83 84 85
  template <typename DeviceContext, typename X, typename Y, typename DX,
            typename DY, typename Dim>
  void operator()(const DeviceContext& place, X& x, Y& y, DX& dx, DY& dy,
G
guosheng 已提交
86 87 88 89
                  const Dim& dim, int size) {
    auto equals = x == y.broadcast(dim);
    auto ones = dx.constant(1);
    auto zeros = dx.constant(0);
90 91
    // If there are multiple minimum or maximum elements, the subgradient of
    // each is the set [0, 1], and we pass gradient to all of them here.
G
guosheng 已提交
92
    dx.device(place) = dy.broadcast(dim) * equals.select(ones, zeros);
G
guosheng 已提交
93 94 95
  }
};

Q
QI JUN 已提交
96
template <typename DeviceContext, typename T, typename Functor>
Y
Yu Yang 已提交
97
class ReduceKernel : public framework::OpKernel<T> {
G
guosheng 已提交
98 99
 public:
  void Compute(const framework::ExecutionContext& context) const override {
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
    bool reduce_all = context.Attr<bool>("reduce_all");
    if (reduce_all) {
      // Flatten and reduce 1-D tensor
      auto* input = context.Input<Tensor>("X");
      auto* output = context.Output<Tensor>("Out");
      output->mutable_data<T>(context.GetPlace());
      auto x = EigenVector<T>::Flatten(*input);
      auto out = EigenScalar<T>::From(*output);
      auto& place =
          *context.template device_context<DeviceContext>().eigen_device();
      auto reduce_dim = Eigen::array<int, 1>({{0}});
      Functor functor;
      functor(place, x, out, reduce_dim);
    } else {
      int rank = context.Input<Tensor>("X")->dims().size();
      switch (rank) {
        case 1:
          ReduceCompute<1>(context);
          break;
        case 2:
          ReduceCompute<2>(context);
          break;
        case 3:
          ReduceCompute<3>(context);
          break;
        case 4:
          ReduceCompute<4>(context);
          break;
        case 5:
          ReduceCompute<5>(context);
          break;
        case 6:
          ReduceCompute<6>(context);
          break;
      }
G
guosheng 已提交
135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
    }
  }

 private:
  template <size_t D>
  void ReduceCompute(const framework::ExecutionContext& context) const {
    auto* input = context.Input<Tensor>("X");
    auto* output = context.Output<Tensor>("Out");
    output->mutable_data<T>(context.GetPlace());

    auto x = EigenTensor<T, D>::From(*input);
    auto x_rank = static_cast<int>(x.dimensions().size());
    int dim = static_cast<int>(context.Attr<int>("dim"));
    if (dim < 0) dim = x_rank + dim;
    auto reduce_dim = Eigen::array<int, 1>({{dim}});
    // construct the squeezed output tensor
G
guosheng 已提交
151
    bool keep_dim = context.Attr<bool>("keep_dim");
G
guosheng 已提交
152 153 154 155 156 157
    DDim dims = output->dims();
    auto dims_vector = vectorize(dims);
    if (keep_dim && x_rank > 1) {
      dims_vector.erase(dims_vector.begin() + dim);
      dims = framework::make_ddim(dims_vector);
    }
D
Dong Zhihong 已提交
158

Q
QI JUN 已提交
159 160
    auto& place =
        *context.template device_context<DeviceContext>().eigen_device();
G
guosheng 已提交
161
    Functor functor;
D
Dong Zhihong 已提交
162 163 164 165 166 167 168 169

    if (D == 1) {
      auto out = EigenScalar<T>::From(*output);
      functor(place, x, out, reduce_dim);
    } else {
      auto out = EigenTensor<T, (D - 1)>::From(*output, dims);
      functor(place, x, out, reduce_dim);
    }
G
guosheng 已提交
170 171 172
  }
};

Q
QI JUN 已提交
173
template <typename DeviceContext, typename T, typename Functor>
Y
Yu Yang 已提交
174
class ReduceGradKernel : public framework::OpKernel<T> {
G
guosheng 已提交
175 176
 public:
  void Compute(const framework::ExecutionContext& context) const override {
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216
    bool reduce_all = context.Attr<bool>("reduce_all");
    if (reduce_all) {
      auto* input0 = context.Input<Tensor>("X");
      auto* input1 = context.Input<Tensor>("Out");
      auto* input2 = context.Input<Tensor>(framework::GradVarName("Out"));
      auto* output = context.Output<Tensor>(framework::GradVarName("X"));
      output->mutable_data<T>(context.GetPlace());
      auto x = EigenVector<T>::Flatten(*input0);
      auto x_reduce = EigenVector<T>::From(*input1);
      auto x_reduce_grad = EigenVector<T>::From(*input2);
      auto x_grad = EigenVector<T>::Flatten(*output);
      auto& place =
          *context.template device_context<DeviceContext>().eigen_device();
      auto broadcast_dim =
          Eigen::array<int, 1>({{static_cast<int>(input0->numel())}});
      Functor functor;
      functor(place, x, x_reduce, x_grad, x_reduce_grad, broadcast_dim,
              broadcast_dim[0]);
    } else {
      int rank = context.Input<Tensor>("X")->dims().size();
      switch (rank) {
        case 1:
          ReduceGradCompute<1>(context);
          break;
        case 2:
          ReduceGradCompute<2>(context);
          break;
        case 3:
          ReduceGradCompute<3>(context);
          break;
        case 4:
          ReduceGradCompute<4>(context);
          break;
        case 5:
          ReduceGradCompute<5>(context);
          break;
        case 6:
          ReduceGradCompute<6>(context);
          break;
      }
G
guosheng 已提交
217 218 219 220 221
    }
  }

 private:
  template <size_t D>
222
  void ReduceGradCompute(const framework::ExecutionContext& context) const {
G
guosheng 已提交
223 224 225 226 227
    auto* input0 = context.Input<Tensor>("X");
    auto* input1 = context.Input<Tensor>("Out");
    auto* input2 = context.Input<Tensor>(framework::GradVarName("Out"));
    auto* output = context.Output<Tensor>(framework::GradVarName("X"));

228 229 230 231 232 233 234 235 236 237 238
    output->mutable_data<T>(context.GetPlace());
    auto x = EigenTensor<T, D>::From(*input0);
    auto x_grad = EigenTensor<T, D>::From(*output);
    auto x_rank = static_cast<int>(x.dimensions().size());
    int dim = static_cast<int>(context.Attr<int>("dim"));
    if (dim < 0) dim = x_rank + dim;
    DDim dims = input0->dims();
    dims[dim] = 1;
    auto x_reduce = EigenTensor<T, D>::From(*input1, dims);
    auto x_reduce_grad = EigenTensor<T, D>::From(*input2, dims);

D
Dong Zhihong 已提交
239 240 241
    Eigen::array<int, D> broadcast_dim;
    for (size_t i = 0; i < D; ++i) broadcast_dim[i] = 1;
    broadcast_dim[dim] = input0->dims()[dim];
Q
QI JUN 已提交
242 243
    auto& place =
        *context.template device_context<DeviceContext>().eigen_device();
244
    Functor functor;
D
Dong Zhihong 已提交
245 246
    functor(place, x, x_reduce, x_grad, x_reduce_grad, broadcast_dim,
            broadcast_dim[dim]);
G
guosheng 已提交
247 248 249 250 251
  }
};

}  // namespace operators
}  // namespace paddle
252 253 254 255 256 257

#define FOR_EACH_KERNEL_FUNCTOR(__macro)                \
  __macro(reduce_sum, SumFunctor, SumGradFunctor);      \
  __macro(reduce_mean, MeanFunctor, MeanGradFunctor);   \
  __macro(reduce_max, MaxFunctor, MaxOrMinGradFunctor); \
  __macro(reduce_min, MinFunctor, MaxOrMinGradFunctor);