pool_with_index_op.cc 9.6 KB
Newer Older
C
chengduoZH 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/* 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. */

#include "paddle/operators/pool_with_index_op.h"

namespace paddle {
namespace operators {

C
chengduoZH 已提交
20 21
inline int OutputSizeMaxPool(int input_size, int filter_size, int padding,
                             int stride) {
C
chengduoZH 已提交
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36
  int output_size = (input_size - filter_size + 2 * padding) / stride + 1;
  return output_size;
}

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

 protected:
  void InferShape(framework::InferShapeContextBase *ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "X(Input) of Pooling should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Out(Output) of Pooling should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("Mask"),
C
chengduoZH 已提交
37
                   "Mask(Output) of Pooling should not be null.");
C
chengduoZH 已提交
38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54

    auto in_x_dims = ctx->GetInputDim("X");

    std::vector<int> ksize = ctx->Attrs().Get<std::vector<int>>("ksize");
    std::vector<int> strides = ctx->Attrs().Get<std::vector<int>>("strides");
    std::vector<int> paddings = ctx->Attrs().Get<std::vector<int>>("paddings");

    PADDLE_ENFORCE(in_x_dims.size() == 4 || in_x_dims.size() == 5,
                   "Pooling intput should be 4-D or 5-D");

    if (ctx->Attrs().Get<bool>("globalPooling")) {
      ksize.resize(static_cast<size_t>(in_x_dims.size()) - 2);
      for (size_t i = 0; i < ksize.size(); ++i)
        ksize[i] = static_cast<int>(in_x_dims[i + 2]);
    }

    PADDLE_ENFORCE(in_x_dims.size() - ksize.size() == 2U,
C
chengduoZH 已提交
55
                   "Intput size and pooling size should be consistent.");
C
chengduoZH 已提交
56
    PADDLE_ENFORCE_EQ(ksize.size(), strides.size(),
C
chengduoZH 已提交
57
                      "Strides size and pooling size should be the same.");
C
chengduoZH 已提交
58
    PADDLE_ENFORCE_EQ(ksize.size(), paddings.size(),
C
chengduoZH 已提交
59
                      "Paddings size and pooling size should be the same.");
C
chengduoZH 已提交
60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76

    std::vector<int64_t> output_shape({in_x_dims[0], in_x_dims[1]});
    for (size_t i = 0; i < ksize.size(); ++i) {
      output_shape.push_back(OutputSizeMaxPool(in_x_dims[i + 2], ksize[i],
                                               paddings[i], strides[i]));
    }
    ctx->SetOutputDim("Out", framework::make_ddim(output_shape));
    ctx->SetOutputDim("Mask", framework::make_ddim(output_shape));
  }
};

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

 protected:
  void InferShape(framework::InferShapeContextBase *ctx) const override {
C
chengduoZH 已提交
77 78 79
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) must not be null.");
    PADDLE_ENFORCE(ctx->HasOutput(framework::GradVarName("X")),
                   "Input(X@GRAD) should not be null.");
C
chengduoZH 已提交
80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
  }
};

class MaxPool2dWithIndexOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  MaxPool2dWithIndexOpMaker(framework::OpProto *proto,
                            framework::OpAttrChecker *op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput(
        "X",
        "The input tensor of pooling operator. "
        "The format of input tensor is NCHW. Where N is batch size, C is the "
        "number of channels, H and W is the height and width of image.");
    AddOutput("Out",
              "The output tensor of pooling operator."
C
chengduoZH 已提交
96 97 98 99
              "The format of output tensor is also NCHW."
              "Where N is batch size, C is "
              "the number of channels, H and W is the height and "
              "width of image.");
C
chengduoZH 已提交
100 101
    AddOutput("Mask",
              "The Mask tensor of pooling operator."
C
chengduoZH 已提交
102 103 104 105
              "The format of output tensor is also NCHW."
              "Where N is batch size, C is the number of channels, H and W "
              "is the height and width of image."
              "The value in it is the index in current feature map");
C
chengduoZH 已提交
106 107

    AddAttr<std::vector<int>>(
C
chengduoZH 已提交
108
        "ksize",
C
chengduoZH 已提交
109
        "The pooling size(height, width) of pooling operator."
C
chengduoZH 已提交
110
        "If globalPooling = true, ksize is ignored and need not be "
C
chengduoZH 已提交
111 112
        "specified.");  // TODO(Chengduo): Add checker. (Currently,
                        // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
113 114
    AddAttr<bool>(
        "globalPooling",
C
chengduoZH 已提交
115 116 117
        "Whether to use the globalPooling."
        "Bool constant equal to false or true."
        "Default false."
C
chengduoZH 已提交
118 119 120
        "If globalPooling = true, ksize is ignored and need not be specified.")
        .SetDefault(false);
    AddAttr<std::vector<int>>("strides",
C
chengduoZH 已提交
121 122
                              "Strides(height, width) of pooling operator."
                              "Default {1,1}.")
C
chengduoZH 已提交
123 124
        .SetDefault({1, 1});  // TODO(Chengduo): Add checker. (Currently,
                              // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
125
    AddAttr<std::vector<int>>("paddings",
C
chengduoZH 已提交
126 127
                              "Paddings(height, width) of pooling operator."
                              "Default {0,0}.")
C
chengduoZH 已提交
128 129
        .SetDefault({0, 0});  // TODO(Chengduo): Add checker. (Currently,
                              // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
130 131

    AddComment(R"DOC(
C
chengduoZH 已提交
132 133 134 135 136 137
The maxPooling2d with index operation calculates the output and the mask
based on the input and ksize, strides, paddings parameters. Input(X) and
output(Out, Mask) are in NCHW format. Where N is batch size, C is the
number of channels, H and W is the height and width of feature.
Parameters(ksize, strides, paddings) are two elements.
These two elements represent height and width, respectively.
C
chengduoZH 已提交
138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154
)DOC");
  }
};

class MaxPool3dWithIndexOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  MaxPool3dWithIndexOpMaker(framework::OpProto *proto,
                            framework::OpAttrChecker *op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput(
        "X",
        "The input tensor of pooling operator. "
        "The format of input tensor is NCDHW. Where N is batch size, C is "
        "the number of channels, D, H and W is the depth, height and width of "
        "image.");
    AddOutput("Out",
              "The output tensor of pooling operator."
C
chengduoZH 已提交
155 156 157 158
              "The format of output tensor is also NCDHW."
              "Where N is batch size, C is "
              "the number of channels, D, H and W is the depth, height and "
              "width of image.");
C
chengduoZH 已提交
159 160
    AddOutput("Mask",
              "The Mask tensor of pooling operator."
C
chengduoZH 已提交
161 162 163 164
              "The format of output tensor is also NCDHW."
              "Where N is batch size, C is the number of channels, D, H and W "
              "is the depth, height and width of image."
              "The value in it is the index in current feature map");
C
chengduoZH 已提交
165 166

    AddAttr<std::vector<int>>(
C
chengduoZH 已提交
167
        "ksize",
C
chengduoZH 已提交
168
        "The pooling size(depth, height, width) of pooling operator."
C
chengduoZH 已提交
169
        "If globalPooling = true, ksize is ignored and need not be "
C
chengduoZH 已提交
170 171
        "specified.");  // TODO(Chengduo): Add checker. (Currently,
                        // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
172 173
    AddAttr<bool>(
        "globalPooling",
C
chengduoZH 已提交
174 175 176
        "Whether to use the globalPooling."
        "Bool constant equal to false or true."
        "Default false."
C
chengduoZH 已提交
177 178 179 180
        "If globalPooling = true, ksize is ignored and need not be specified.")
        .SetDefault(false);
    AddAttr<std::vector<int>>(
        "strides",
C
chengduoZH 已提交
181 182
        "Strides(depth, height, width) of pooling operator."
        "Default {1,1,1}.")
C
chengduoZH 已提交
183 184
        .SetDefault({1, 1, 1});  // TODO(Chengduo): Add checker. (Currently,
                                 // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
185 186
    AddAttr<std::vector<int>>(
        "paddings",
C
chengduoZH 已提交
187 188
        "Paddings(depth, height, width) of pooling operator."
        "Default {0,0,0}.")
C
chengduoZH 已提交
189 190
        .SetDefault({0, 0, 0});  // TODO(Chengduo): Add checker. (Currently,
                                 // TypedAttrChecker don't support vector type.)
C
chengduoZH 已提交
191

C
chengduoZH 已提交
192
    AddComment(R"DOC(
C
chengduoZH 已提交
193 194 195 196 197 198
The maxpooling3d with index operation calculates the output and the mask
based on the input and ksize, strides, paddings parameters.
Input(X) and output(Out, Mask) are in NCDHW format. Where N is batch
size, C is the number of channels, D, H and W is the depth, height and
width of feature. Parameters(ksize, strides, paddings) are three elements.
These three elements represent depth, height and width, respectively.
C
chengduoZH 已提交
199 200 201
)DOC");
  }
};
C
chengduoZH 已提交
202

C
chengduoZH 已提交
203 204 205 206 207
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

C
chengduoZH 已提交
208 209
REGISTER_OP(max_pool2d_with_index, ops::MaxPoolWithIndexOp,
            ops::MaxPool2dWithIndexOpMaker, max_pool2d_with_index_grad,
C
chengduoZH 已提交
210 211 212
            ops::MaxPoolWithIndexOpGrad);

REGISTER_OP_CPU_KERNEL(
C
chengduoZH 已提交
213
    max_pool2d_with_index,
C
chengduoZH 已提交
214 215
    ops::MaxPoolWithIndexKernel<paddle::platform::CPUPlace, float>);
REGISTER_OP_CPU_KERNEL(
C
chengduoZH 已提交
216
    max_pool2d_with_index_grad,
C
chengduoZH 已提交
217 218
    ops::MaxPoolWithIndexGradKernel<paddle::platform::CPUPlace, float>)

C
chengduoZH 已提交
219 220
REGISTER_OP(max_pool3d_with_index, ops::MaxPoolWithIndexOp,
            ops::MaxPool3dWithIndexOpMaker, max_pool3d_with_index_grad,
C
chengduoZH 已提交
221 222 223
            ops::MaxPoolWithIndexOpGrad);

REGISTER_OP_CPU_KERNEL(
C
chengduoZH 已提交
224
    max_pool3d_with_index,
C
chengduoZH 已提交
225 226
    ops::MaxPoolWithIndexKernel<paddle::platform::CPUPlace, float>);
REGISTER_OP_CPU_KERNEL(
C
chengduoZH 已提交
227
    max_pool3d_with_index_grad,
C
chengduoZH 已提交
228
    ops::MaxPoolWithIndexGradKernel<paddle::platform::CPUPlace, float>)