reduce_op.cc 7.9 KB
Newer Older
G
guosheng 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

L
Luo Tao 已提交
3 4 5
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
G
guosheng 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
G
guosheng 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
G
guosheng 已提交
14 15 16 17 18 19 20 21 22 23 24 25

#include "paddle/operators/reduce_op.h"

namespace paddle {
namespace operators {

using framework::Tensor;

class ReduceOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

26
  void InferShape(framework::InferShapeContext *ctx) const override {
27 28 29 30 31
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "Input(X) of ReduceOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of ReduceOp should not be null.");
    auto x_dims = ctx->GetInputDim("X");
G
guosheng 已提交
32
    auto x_rank = x_dims.size();
G
guosheng 已提交
33
    PADDLE_ENFORCE_LE(x_rank, 6, "Tensors with rank at most 6 are supported.");
34
    int dim = ctx->Attrs().Get<int>("dim");
G
guosheng 已提交
35 36 37
    if (dim < 0) dim = x_rank + dim;
    PADDLE_ENFORCE_LT(
        dim, x_rank,
G
guosheng 已提交
38
        "The dim should be in the range [-rank(input), rank(input)).");
39
    bool reduce_all = ctx->Attrs().Get<bool>("reduce_all");
40
    bool keep_dim = ctx->Attrs().Get<bool>("keep_dim");
41
    if (reduce_all) {
42 43 44 45 46
      if (keep_dim)
        ctx->SetOutputDim(
            "Out", framework::make_ddim(std::vector<int64_t>(x_rank, 1)));
      else
        ctx->SetOutputDim("Out", {1});
G
guosheng 已提交
47
    } else {
48 49 50 51 52 53 54 55 56 57 58 59
      auto dims_vector = vectorize(x_dims);
      if (keep_dim || x_rank == 1) {
        dims_vector[dim] = 1;
      } else {
        dims_vector.erase(dims_vector.begin() + dim);
      }
      auto out_dims = framework::make_ddim(dims_vector);
      ctx->SetOutputDim("Out", out_dims);
      if (dim != 0) {
        // Only pass LoD when not reducing on the first dim.
        ctx->ShareLoD("X", /*->*/ "Out");
      }
60
    }
G
guosheng 已提交
61 62 63 64 65 66 67
  }
};

class ReduceGradOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

68
  void InferShape(framework::InferShapeContext *ctx) const override {
69 70 71 72
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null.");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null.");
    auto x_dims = ctx->GetInputDim("X");
G
guosheng 已提交
73
    auto x_rank = x_dims.size();
G
guosheng 已提交
74
    PADDLE_ENFORCE_LE(x_rank, 6, "Tensors with rank at most 6 are supported.");
75
    int dim = ctx->Attrs().Get<int>("dim");
G
guosheng 已提交
76 77 78
    if (dim < 0) dim = x_rank + dim;
    PADDLE_ENFORCE_LT(
        dim, x_rank,
G
guosheng 已提交
79
        "The dim should be in the range [-rank(input), rank(input)).");
80 81 82
    auto x_grad_name = framework::GradVarName("X");
    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
83
      ctx->ShareLoD("X", /*->*/ x_grad_name);
84
    }
G
guosheng 已提交
85 86 87
  }
};

G
guosheng 已提交
88
class ReduceOpMaker : public framework::OpProtoAndCheckerMaker {
G
guosheng 已提交
89
 public:
90
  ReduceOpMaker(OpProto *proto, OpAttrChecker *op_checker)
G
guosheng 已提交
91
      : OpProtoAndCheckerMaker(proto, op_checker) {
K
kexinzhao 已提交
92 93 94
    AddInput("X",
             "(Tensor) The input tensor. Tensors with rank at most 6 are "
             "supported.");
G
guosheng 已提交
95
    AddOutput("Out", "(Tensor) The result tensor.");
96 97
    AddAttr<int>(
        "dim",
K
kexinzhao 已提交
98
        "(int, default 0) The dimension to reduce. "
99 100
        "Must be in the range [-rank(input), rank(input)). "
        "If `dim < 0`, the dim to reduce is `rank + dim`. "
K
kexinzhao 已提交
101
        "Note that reducing on the first dim will make the LoD info lost.")
102
        .SetDefault(0);
G
guosheng 已提交
103 104 105 106
    AddAttr<bool>("keep_dim",
                  "(bool, default false) "
                  "If true, retain the reduced dimension with length 1.")
        .SetDefault(false);
107 108 109 110
    AddAttr<bool>("reduce_all",
                  "(bool, default false) "
                  "If true, output a scalar reduced along all dimensions.")
        .SetDefault(false);
G
guosheng 已提交
111
    comment_ = R"DOC(
K
kexinzhao 已提交
112 113 114 115
{ReduceOp} Operator.

This operator computes the {reduce} of input tensor along the given dimension. 
The result tensor has 1 fewer dimension than the input unless keep_dim is true.
116
If reduce_all is true, just reduce along all dimensions and output a scalar.
K
kexinzhao 已提交
117

G
guosheng 已提交
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
)DOC";
    AddComment(comment_);
  }

 protected:
  std::string comment_;

  void Replace(std::string &src, std::string from, std::string to) {
    std::size_t len_from = std::strlen(from.c_str());
    std::size_t len_to = std::strlen(to.c_str());
    for (std::size_t pos = src.find(from); pos != std::string::npos;
         pos = src.find(from, pos + len_to)) {
      src.replace(pos, len_from, to);
    }
  }

  void SetComment(std::string name, std::string op) {
135
    Replace(comment_, "{ReduceOp}", name);
G
guosheng 已提交
136
    Replace(comment_, "{reduce}", op);
G
guosheng 已提交
137 138 139
  }
};

G
guosheng 已提交
140 141
class ReduceSumOpMaker : public ReduceOpMaker {
 public:
142
  ReduceSumOpMaker(OpProto *proto, OpAttrChecker *op_checker)
G
guosheng 已提交
143 144 145 146 147 148 149
      : ReduceOpMaker(proto, op_checker) {
    SetComment("ReduceSum", "sum");
    AddComment(comment_);
  }
};

class ReduceMeanOpMaker : public ReduceOpMaker {
G
guosheng 已提交
150
 public:
151
  ReduceMeanOpMaker(OpProto *proto, OpAttrChecker *op_checker)
G
guosheng 已提交
152 153 154
      : ReduceOpMaker(proto, op_checker) {
    SetComment("ReduceMean", "mean");
    AddComment(comment_);
G
guosheng 已提交
155 156 157
  }
};

G
guosheng 已提交
158
class ReduceMaxOpMaker : public ReduceOpMaker {
G
guosheng 已提交
159
 public:
160
  ReduceMaxOpMaker(OpProto *proto, OpAttrChecker *op_checker)
G
guosheng 已提交
161 162 163
      : ReduceOpMaker(proto, op_checker) {
    SetComment("ReduceMax", "max");
    AddComment(comment_);
G
guosheng 已提交
164 165 166
  }
};

G
guosheng 已提交
167
class ReduceMinOpMaker : public ReduceOpMaker {
G
guosheng 已提交
168
 public:
169
  ReduceMinOpMaker(OpProto *proto, OpAttrChecker *op_checker)
G
guosheng 已提交
170 171 172
      : ReduceOpMaker(proto, op_checker) {
    SetComment("ReduceMin", "min");
    AddComment(comment_);
G
guosheng 已提交
173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP(reduce_sum, ops::ReduceOp, ops::ReduceSumOpMaker, reduce_sum_grad,
            ops::ReduceGradOp);

REGISTER_OP(reduce_mean, ops::ReduceOp, ops::ReduceMeanOpMaker,
            reduce_mean_grad, ops::ReduceGradOp);

REGISTER_OP(reduce_max, ops::ReduceOp, ops::ReduceMaxOpMaker, reduce_max_grad,
            ops::ReduceGradOp);

L
Luo Tao 已提交
190
REGISTER_OP(reduce_min, ops::ReduceOp, ops::ReduceMinOpMaker, reduce_min_grad,
G
guosheng 已提交
191
            ops::ReduceGradOp);
G
guosheng 已提交
192

Q
QI JUN 已提交
193 194 195
#define REGISTER_REDUCE_CPU_KERNEL(reduce_type, functor, grad_functor)         \
  REGISTER_OP_CPU_KERNEL(reduce_type,                                          \
                         ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
196 197 198 199 200 201 202
                                           float, ops::functor>,               \
                         ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
                                           double, ops::functor>,              \
                         ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
                                           int, ops::functor>,                 \
                         ops::ReduceKernel<paddle::platform::CPUDeviceContext, \
                                           int64_t, ops::functor>);            \
Q
QI JUN 已提交
203 204 205
  REGISTER_OP_CPU_KERNEL(                                                      \
      reduce_type##_grad,                                                      \
      ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, float,         \
206 207 208 209 210 211
                            ops::grad_functor>,                                \
      ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, double,        \
                            ops::grad_functor>,                                \
      ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, int,           \
                            ops::grad_functor>,                                \
      ops::ReduceGradKernel<paddle::platform::CPUDeviceContext, int64_t,       \
Q
QI JUN 已提交
212
                            ops::grad_functor>);
213 214

FOR_EACH_KERNEL_FUNCTOR(REGISTER_REDUCE_CPU_KERNEL);