reduce_op.h 6.4 KB
Newer Older
G
guosheng 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
/* 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

#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>;

struct SumFunctor {
G
guosheng 已提交
30 31 32
  template <typename Place, typename X, typename Y, typename Dim>
  void operator()(const Place& place, X& x, Y& y, const Dim& dim) {
    y.device(place) = x.sum(dim);
G
guosheng 已提交
33 34 35 36
  }
};

struct SumGradFunctor {
G
guosheng 已提交
37
  template <typename Place, typename X, typename Y, typename DX, typename DY,
G
guosheng 已提交
38
            typename Dim>
G
guosheng 已提交
39 40 41
  void operator()(const Place& place, X& x, Y& y, DX& dx, DY& dy,
                  const Dim& dim, int size) {
    dx.device(place) = dy.broadcast(dim);
G
guosheng 已提交
42 43 44 45
  }
};

struct MeanFunctor {
G
guosheng 已提交
46 47 48
  template <typename Place, typename X, typename Y, typename Dim>
  void operator()(const Place& place, X& x, Y& y, const Dim& dim) {
    y.device(place) = x.mean(dim);
G
guosheng 已提交
49 50 51 52
  }
};

struct MeanGradFunctor {
G
guosheng 已提交
53
  template <typename Place, typename X, typename Y, typename DX, typename DY,
G
guosheng 已提交
54
            typename Dim>
G
guosheng 已提交
55 56 57
  void operator()(const Place& place, X& x, Y& y, DX& dx, DY& dy,
                  const Dim& dim, int size) {
    dx.device(place) = dy.broadcast(dim) / dx.constant(size);
G
guosheng 已提交
58 59 60 61
  }
};

struct MaxFunctor {
G
guosheng 已提交
62 63 64
  template <typename Place, typename X, typename Y, typename Dim>
  void operator()(const Place& place, X& x, Y& y, const Dim& dim) {
    y.device(place) = x.maximum(dim);
G
guosheng 已提交
65 66 67 68
  }
};

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

struct MaxOrMinGradFunctor {
G
guosheng 已提交
76
  template <typename Place, typename X, typename Y, typename DX, typename DY,
G
guosheng 已提交
77
            typename Dim>
G
guosheng 已提交
78 79 80 81 82
  void operator()(const Place& place, X& x, Y& y, DX& dx, DY& dy,
                  const Dim& dim, int size) {
    auto equals = x == y.broadcast(dim);
    auto ones = dx.constant(1);
    auto zeros = dx.constant(0);
83 84
    // 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 已提交
85
    dx.device(place) = dy.broadcast(dim) * equals.select(ones, zeros);
G
guosheng 已提交
86 87 88 89 90 91 92 93 94 95 96 97 98 99 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
  }
};

template <typename Place, typename T, typename Functor>
class ReduceKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    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;
    }
  }

 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 已提交
129
    bool keep_dim = context.Attr<bool>("keep_dim");
G
guosheng 已提交
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149
    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);
    }
    auto out = EigenTensor < T, D == 1 ? 1 : (D - 1) > ::From(*output, dims);
    auto& place = context.GetEigenDevice<Place>();
    Functor functor;
    functor(place, x, out, reduce_dim);
  }
};

template <typename Place, typename T, typename Functor>
class ReduceGradKernel : public framework::OpKernel {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
    int rank = context.Input<Tensor>("X")->dims().size();
    switch (rank) {
      case 1:
150
        ReduceGradCompute<1>(context);
G
guosheng 已提交
151 152
        break;
      case 2:
153
        ReduceGradCompute<2>(context);
G
guosheng 已提交
154 155
        break;
      case 3:
156
        ReduceGradCompute<3>(context);
G
guosheng 已提交
157 158
        break;
      case 4:
159
        ReduceGradCompute<4>(context);
G
guosheng 已提交
160 161
        break;
      case 5:
162
        ReduceGradCompute<5>(context);
G
guosheng 已提交
163 164
        break;
      case 6:
165
        ReduceGradCompute<6>(context);
G
guosheng 已提交
166 167 168 169 170 171
        break;
    }
  }

 private:
  template <size_t D>
172
  void ReduceGradCompute(const framework::ExecutionContext& context) const {
G
guosheng 已提交
173 174 175 176 177
    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"));

178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
    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);

    Eigen::array<int, D> braodcast_dim;
    for (size_t i = 0; i < D; ++i) braodcast_dim[i] = 1;
    braodcast_dim[dim] = input0->dims()[dim];
    auto& place = context.GetEigenDevice<Place>();
    Functor functor;
    functor(place, x, x_reduce, x_grad, x_reduce_grad, braodcast_dim,
            braodcast_dim[dim]);
G
guosheng 已提交
196 197 198 199 200
  }
};

}  // namespace operators
}  // namespace paddle