pool_op.h 6.8 KB
Newer Older
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
/* 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"
#include "paddle/operators/math/math_function.h"
#include "paddle/operators/math/pooling.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51

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

  void InferShape(framework::InferShapeContext* ctx) const override;
};

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

  void InferShape(framework::InferShapeContext* ctx) const override;
};

class Pool2dOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  Pool2dOpMaker(framework::OpProto* proto,
                framework::OpAttrChecker* op_checker);
};

class Pool3dOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  Pool3dOpMaker(framework::OpProto* proto,
                framework::OpAttrChecker* op_checker);
};
52 53

template <typename Place, typename T>
C
chengduoZH 已提交
54
class PoolKernel : public framework::OpKernel<T> {
55 56
 public:
  void Compute(const framework::ExecutionContext& context) const override {
C
chengduoZH 已提交
57
    const Tensor* in_x = context.Input<Tensor>("X");
58
    Tensor* out = context.Output<Tensor>("Out");
59

C
chengduoZH 已提交
60
    std::string pooling_type = context.Attr<std::string>("pooling_type");
61 62 63
    std::vector<int> ksize = context.Attr<std::vector<int>>("ksize");
    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");
C
chengduoZH 已提交
64
    if (context.Attr<bool>("global_pooling")) {
65
      for (size_t i = 0; i < ksize.size(); ++i) {
C
fix bug  
chengduoZH 已提交
66
        paddings[i] = 0;
C
chengduoZH 已提交
67
        ksize[i] = static_cast<int>(in_x->dims()[i + 2]);
68 69 70 71 72 73
      }
    }

    switch (ksize.size()) {
      case 2: {
        if (pooling_type == "max") {
C
chengduoZH 已提交
74
          paddle::operators::math::Pool2dFunctor<
75
              Place, paddle::operators::math::MaxPool<T>, T>
76
              pool2d_forward;
77
          paddle::operators::math::MaxPool<T> pool_process;
C
chengduoZH 已提交
78
          pool2d_forward(context.device_context(), *in_x, *out, ksize, strides,
79
                         paddings, pool_process);
80

C
chengduoZH 已提交
81
        } else if (pooling_type == "avg") {
C
chengduoZH 已提交
82
          paddle::operators::math::Pool2dFunctor<
83
              Place, paddle::operators::math::AvgPool<T>, T>
84
              pool2d_forward;
85
          paddle::operators::math::AvgPool<T> pool_process;
C
chengduoZH 已提交
86
          pool2d_forward(context.device_context(), *in_x, *out, ksize, strides,
87
                         paddings, pool_process);
88 89 90 91
        }
      } break;
      case 3: {
        if (pooling_type == "max") {
C
chengduoZH 已提交
92
          paddle::operators::math::Pool3dFunctor<
93
              Place, paddle::operators::math::MaxPool<T>, T>
94
              pool3d_forward;
95
          paddle::operators::math::MaxPool<T> pool_process;
C
chengduoZH 已提交
96
          pool3d_forward(context.device_context(), *in_x, *out, ksize, strides,
97
                         paddings, pool_process);
C
chengduoZH 已提交
98
        } else if (pooling_type == "avg") {
C
chengduoZH 已提交
99
          paddle::operators::math::Pool3dFunctor<
100
              Place, paddle::operators::math::AvgPool<T>, T>
101
              pool3d_forward;
102
          paddle::operators::math::AvgPool<T> pool_process;
C
chengduoZH 已提交
103
          pool3d_forward(context.device_context(), *in_x, *out, ksize, strides,
104
                         paddings, pool_process);
105 106
        }
      } break;
C
fix bug  
chengduoZH 已提交
107
      default: { PADDLE_THROW("Pool op only supports 2D and 3D input."); }
108 109 110 111 112
    }
  }
};

template <typename Place, typename T>
C
chengduoZH 已提交
113
class PoolGradKernel : public framework::OpKernel<T> {
114 115
 public:
  void Compute(const framework::ExecutionContext& context) const override {
C
chengduoZH 已提交
116
    const Tensor* in_x = context.Input<Tensor>("X");
117 118 119
    const Tensor* out = context.Input<Tensor>("Out");
    const Tensor* out_grad =
        context.Input<Tensor>(framework::GradVarName("Out"));
C
chengduoZH 已提交
120
    Tensor* in_x_grad = context.Output<Tensor>(framework::GradVarName("X"));
121

C
chengduoZH 已提交
122
    std::string pooling_type = context.Attr<std::string>("pooling_type");
123 124 125 126
    std::vector<int> ksize = context.Attr<std::vector<int>>("ksize");
    std::vector<int> strides = context.Attr<std::vector<int>>("strides");
    std::vector<int> paddings = context.Attr<std::vector<int>>("paddings");

C
chengduoZH 已提交
127
    if (context.Attr<bool>("global_pooling")) {
C
fix bug  
chengduoZH 已提交
128 129
      for (size_t i = 0; i < ksize.size(); ++i) {
        paddings[i] = 0;
C
chengduoZH 已提交
130
        ksize[i] = static_cast<int>(in_x->dims()[i + 2]);
C
fix bug  
chengduoZH 已提交
131
      }
132 133
    }

C
chengduoZH 已提交
134 135 136
    if (in_x_grad) {
      in_x_grad->mutable_data<T>(context.GetPlace());
      auto temp = framework::EigenVector<T>::Flatten(*in_x_grad);
137 138 139 140 141 142
      temp.device(context.GetEigenDevice<Place>()) =
          temp.constant(static_cast<T>(0));

      switch (ksize.size()) {
        case 2: {
          if (pooling_type == "max") {
C
chengduoZH 已提交
143
            paddle::operators::math::MaxPool2dGradFunctor<Place, T>
144
                pool2d_backward;
C
chengduoZH 已提交
145
            pool2d_backward(context.device_context(), *in_x, *in_x_grad, *out,
C
chengduoZH 已提交
146
                            *out_grad, ksize, strides, paddings);
C
chengduoZH 已提交
147
          } else if (pooling_type == "avg") {
C
chengduoZH 已提交
148
            paddle::operators::math::Pool2dGradFunctor<
149
                Place, paddle::operators::math::AvgPoolGrad<T>, T>
150
                pool2d_backward;
151
            paddle::operators::math::AvgPoolGrad<T> pool_process;
C
chengduoZH 已提交
152
            pool2d_backward(context.device_context(), *in_x, *in_x_grad, *out,
153
                            *out_grad, ksize, strides, paddings, pool_process);
154 155 156 157
          }
        } break;
        case 3: {
          if (pooling_type == "max") {
C
chengduoZH 已提交
158
            paddle::operators::math::MaxPool3dGradFunctor<Place, T>
159
                pool3d_backward;
C
chengduoZH 已提交
160
            pool3d_backward(context.device_context(), *in_x, *in_x_grad, *out,
C
chengduoZH 已提交
161
                            *out_grad, ksize, strides, paddings);
C
chengduoZH 已提交
162
          } else if (pooling_type == "avg") {
C
chengduoZH 已提交
163
            paddle::operators::math::Pool3dGradFunctor<
164
                Place, paddle::operators::math::AvgPoolGrad<T>, T>
165
                pool3d_backward;
166
            paddle::operators::math::AvgPoolGrad<T> pool_process;
C
chengduoZH 已提交
167
            pool3d_backward(context.device_context(), *in_x, *in_x_grad, *out,
168
                            *out_grad, ksize, strides, paddings, pool_process);
169 170
          }
        } break;
C
fix bug  
chengduoZH 已提交
171
        default: { PADDLE_THROW("Pool op only supports 2D and 3D input."); }
172 173 174 175 176 177 178
      }
    }
  }
};

}  // namespace operators
}  // namespace paddle