diff --git a/mace/kernels/neon/pooling_neon.cc b/mace/kernels/neon/pooling_neon.cc index 6960c7eefbb2ca26571608b0783f339d4431ea28..7112234deb16b6f899e04ed3186fd4b84c34b5b0 100644 --- a/mace/kernels/neon/pooling_neon.cc +++ b/mace/kernels/neon/pooling_neon.cc @@ -3,8 +3,6 @@ // #include "mace/kernels/pooling.h" -#include -#include "mace/kernels/conv_pool_2d_util.h" namespace mace { namespace kernels { @@ -61,9 +59,15 @@ void PoolingFunctor::operator()( const index_t *input_shape, float *output, const index_t *output_shape) { + + int paddings[2]; + std::vector filter_shape = {input_shape[1], input_shape[0], + kernels_[0], kernels_[1]}; + kernels::CalPaddingSize(input_shape, filter_shape.data(), this->dilations_, + strides_, this->padding_, paddings); #ifdef __COPY_MAKE_PADDING Tensor padded_input; - ConstructInputWithPadding(input, input_shape, paddings_, &padded_input); + ConstructInputWithPadding(input, input_shape, paddings, &padded_input); input = padded_input.data(); input_shape = padded_input.shape().data(); #endif @@ -76,14 +80,14 @@ void PoolingFunctor::operator()( PoolingMaxNeonK2x2S2x2Padded(input, input_shape, output, output_shape); #else PoolingMaxNeonK2x2S2x2(input, input_shape, output, output_shape, - paddings_); + paddings); #endif } else { // AVG_POOL_2x2s2x2 #ifdef __COPY_MAKE_PADDING PoolingAvgNeonK2x2S2x2Padded(input, input_shape, output, output_shape); #else PoolingAvgNeonK2x2S2x2(input, input_shape, output, output_shape, - paddings_); + paddings); #endif } } else if (kernels_[0] == 3 && kernels_[1] == 3 && strides_[0] == 2 && @@ -94,19 +98,19 @@ void PoolingFunctor::operator()( PoolingMaxNeonK3x3S2x2Padded(input, input_shape, output, output_shape); #else PoolingMaxNeonK3x3S2x2(input, input_shape, output, output_shape, - paddings_); + paddings); #endif } else { // AVG_POOL_3x3s2x2 #ifdef __COPY_MAKE_PADDING PoolingAvgNeonK3x3S2x2Padded(input, input_shape, output, output_shape); #else PoolingAvgNeonK3x3S2x2(input, input_shape, output, output_shape, - paddings_); + paddings); #endif } } else { // not implement yet PoolingFunctor(pooling_type_, kernels_, strides_, - paddings_, dilations_)( + padding_, dilations_)( input, input_shape, output, output_shape); } } diff --git a/mace/kernels/pooling.h b/mace/kernels/pooling.h index 1e69cf90b87f9c1e859a7794bc3d3cc88c4d4ee8..821baf3a85e63dbbba5b0ceed9c010dd71dc6071 100644 --- a/mace/kernels/pooling.h +++ b/mace/kernels/pooling.h @@ -7,6 +7,7 @@ #include #include "mace/core/tensor.h" +#include "mace/kernels/conv_pool_2d_util.h" namespace mace { @@ -22,12 +23,12 @@ struct PoolingFunctor { PoolingFunctor(const PoolingType pooling_type, const int *kernels, const int *strides, - const int *paddings, + const Padding padding, const int *dilations) : pooling_type_(pooling_type), kernels_(kernels), strides_(strides), - paddings_(paddings), + padding_(padding), dilations_(dilations) {} void operator()(const T *input, @@ -54,9 +55,14 @@ struct PoolingFunctor { int dilation_h = dilations_[0]; int dilation_w = dilations_[1]; + int paddings[2]; + std::vector filter_shape = {input_shape[1], input_shape[0], + kernels_[0], kernels_[1]}; + kernels::CalPaddingSize(input_shape, filter_shape.data(), this->dilations_, + strides_, this->padding_, paddings); // The left-upper most offset of the padded input - int padded_h_start = 0 - paddings_[0] / 2; - int padded_w_start = 0 - paddings_[1] / 2; + int padded_h_start = 0 - paddings[0] / 2; + int padded_w_start = 0 - paddings[1] / 2; if (pooling_type_ == MAX) { #pragma omp parallel for collapse(2) @@ -116,7 +122,7 @@ struct PoolingFunctor { const PoolingType pooling_type_; const int *kernels_; const int *strides_; - const int *paddings_; + const Padding padding_; const int *dilations_; }; @@ -127,6 +133,13 @@ void PoolingFunctor::operator()( float *output, const index_t *output_shape); +template <> +void PoolingFunctor::operator()( + const float *input, + const index_t *input_shape, + float *output, + const index_t *output_shape); + } // namespace kernels } // namespace mace diff --git a/mace/ops/BUILD b/mace/ops/BUILD index e823136d965670cc198ed14c0f682cbf0b152d00..cacf9103a7b2335b3cd85dd4cc76a5ce6c3dcd06 100644 --- a/mace/ops/BUILD +++ b/mace/ops/BUILD @@ -62,6 +62,22 @@ cc_test( ], ) +cc_test( + name = "pooling_test", + testonly = 1, + srcs = glob( + ["pooling_test.cc"], + ), + copts = ["-std=c++11"], + linkopts = ["-fopenmp"] + if_android(["-ldl"]), + linkstatic = 1, + deps = [ + ":ops", + ":test", + "@gtest//:gtest_main", + ], +) + cc_test( name = "ops_benchmark", testonly = 1, diff --git a/mace/ops/pooling.h b/mace/ops/pooling.h index 15eb120bae51ece80a41a18523e13f025433344f..d1172d6f64f038ad6215dc27c47d21ea5f01fb68 100644 --- a/mace/ops/pooling.h +++ b/mace/ops/pooling.h @@ -19,7 +19,9 @@ class PoolingOp : public ConvPool2dOpBase { kernels_(OperatorBase::GetRepeatedArgument("kernels")), pooling_type_( static_cast(OperatorBase::GetSingleArgument( - "pooling_type", static_cast(AVG)))){}; + "pooling_type", static_cast(AVG)))), + functor_(pooling_type_, kernels_.data(), ConvPool2dOpBase::strides_.data(), + ConvPool2dOpBase::padding_, ConvPool2dOpBase::dilations_.data()){}; bool Run() override { const Tensor *input = this->Input(INPUT); @@ -40,10 +42,7 @@ class PoolingOp : public ConvPool2dOpBase { paddings.data()); output->Resize(output_shape); - auto pooling_func = kernels::PoolingFunctor( - pooling_type_, kernels_.data(), this->strides_.data(), paddings.data(), - this->dilations_.data()); - pooling_func(input->data(), input->shape().data(), + functor_(input->data(), input->shape().data(), output->mutable_data(), output->shape().data()); return true; }; @@ -51,6 +50,7 @@ class PoolingOp : public ConvPool2dOpBase { protected: std::vector kernels_; PoolingType pooling_type_; + kernels::PoolingFunctor functor_; OP_INPUT_TAGS(INPUT); OP_OUTPUT_TAGS(OUTPUT);