This is an automated email from the ASF dual-hosted git repository.
zhasheng pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/v1.x by this push:
new 6cb49f0 [v1.x] Merge MKLDNN adaptive pooling with traditional pooling
implementation (#20699)
6cb49f0 is described below
commit 6cb49f09c42a533ffc395175e9199dadf3268f4d
Author: PiotrWolinski - Intel <[email protected]>
AuthorDate: Thu Nov 4 21:32:31 2021 +0100
[v1.x] Merge MKLDNN adaptive pooling with traditional pooling
implementation (#20699)
---
src/operator/contrib/adaptive_avg_pooling-inl.h | 54 +++----
src/operator/contrib/adaptive_avg_pooling.cc | 19 +--
.../nn/mkldnn/mkldnn_adaptive_pooling-inl.h | 155 ---------------------
src/operator/nn/mkldnn/mkldnn_adaptive_pooling.cc | 100 -------------
src/operator/nn/mkldnn/mkldnn_pooling-inl.h | 26 +++-
src/operator/nn/mkldnn/mkldnn_pooling.cc | 41 ++++--
src/operator/nn/pooling-inl.h | 9 +-
src/operator/nn/pooling.cc | 5 +-
.../mkldnn/mkldnn_quantized_pooling.cc | 3 +-
9 files changed, 95 insertions(+), 317 deletions(-)
diff --git a/src/operator/contrib/adaptive_avg_pooling-inl.h
b/src/operator/contrib/adaptive_avg_pooling-inl.h
index e713de6..118200a 100644
--- a/src/operator/contrib/adaptive_avg_pooling-inl.h
+++ b/src/operator/contrib/adaptive_avg_pooling-inl.h
@@ -43,24 +43,11 @@
#include "../operator_common.h"
#include "../mxnet_op.h"
#include "../mshadow_op.h"
-#if MXNET_USE_MKLDNN == 1
-#include "../nn/mkldnn/mkldnn_adaptive_pooling-inl.h"
-#endif
+#include "../nn/pooling-inl.h"
namespace mxnet {
namespace op {
-struct AdaptiveAvgPoolParam : public dmlc::Parameter<AdaptiveAvgPoolParam> {
- mxnet::Tuple<int> output_size;
- DMLC_DECLARE_PARAMETER(AdaptiveAvgPoolParam) {
- DMLC_DECLARE_FIELD(output_size).set_default(mxnet::Tuple<int>())
- .describe("int (output size) or a tuple of int for output (height,
width).");
- }
- bool operator==(const AdaptiveAvgPoolParam &other) const {
- return this->output_size == other.output_size;
- }
-};
-
static inline bool IsWriting(const OpReqType ort) {
return ort == kWriteTo || ort == kWriteInplace;
}
@@ -128,18 +115,25 @@ static bool AdaptiveAvgPoolOpInferShape(const
nnvm::NodeAttrs& attrs,
using namespace mshadow;
CHECK_EQ(in_shape->size(), 1U) << "Input:[data]";
CHECK_EQ(out_shape->size(), 1U) << "Output:[data]";
- const AdaptiveAvgPoolParam& param =
nnvm::get<AdaptiveAvgPoolParam>(attrs.parsed);
+ const PoolingParam& param = nnvm::get<PoolingParam>(attrs.parsed);
+
mxnet::TShape dshape(in_shape->at(0));
- if (mxnet::op::shape_is_none(dshape)) return false;
- if (param.output_size.ndim() == 0) {
- dshape[2] = 1;
- dshape[3] = 1;
- } else if (param.output_size.ndim() == 1) {
- dshape[2] = param.output_size[0];
- dshape[3] = param.output_size[0];
- } else if (param.output_size.ndim() == 2) {
- dshape[2] = param.output_size[0];
- dshape[3] = param.output_size[1];
+
+ if (mxnet::op::shape_is_none(dshape)) {
+ return false;
+ }
+
+ if (param.output_size.has_value()) {
+ if (param.output_size.value().ndim() == 1) {
+ dshape[2] = param.output_size.value()[0];
+ dshape[3] = param.output_size.value()[0];
+ } else if (param.output_size.value().ndim() == 2) {
+ dshape[2] = param.output_size.value()[0];
+ dshape[3] = param.output_size.value()[1];
+ } else {
+ dshape[2] = 1;
+ dshape[3] = 1;
+ }
} else {
dshape[2] = 1;
dshape[3] = 1;
@@ -161,14 +155,4 @@ MSHADOW_XINLINE int get_stride(Tensor<xpu, Dim, DType>
tensor, int idx) {
} // namespace op
} // namespace mxnet
-namespace std {
-template <>
-struct hash<mxnet::op::AdaptiveAvgPoolParam> {
- size_t operator()(const mxnet::op::AdaptiveAvgPoolParam &val) {
- size_t ret = 0;
- ret = dmlc::HashCombine(ret, val.output_size);
- return ret;
- }
-};
-} // namespace std
#endif // MXNET_OPERATOR_CONTRIB_ADAPTIVE_AVG_POOLING_INL_H_
diff --git a/src/operator/contrib/adaptive_avg_pooling.cc
b/src/operator/contrib/adaptive_avg_pooling.cc
index 56abfe7..a11f88c 100644
--- a/src/operator/contrib/adaptive_avg_pooling.cc
+++ b/src/operator/contrib/adaptive_avg_pooling.cc
@@ -26,6 +26,9 @@
#include "../elemwise_op_common.h"
#include "../operator_common.h"
#include "adaptive_avg_pooling-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "../nn/mkldnn/mkldnn_pooling-inl.h"
+#endif // MXNET_USE_MKLDNN
#define START_IND(a, b, c) static_cast<int>(std::floor(static_cast<float>(a *
c) / b))
#define END_IND(a, b, c) static_cast<int>(std::ceil(static_cast<float>((a + 1)
* c) / b))
@@ -240,12 +243,12 @@ void AdaptiveAvgPoolComputeExCPU(const nnvm::NodeAttrs
&attrs,
*/
if (SupportMKLDNN(inputs[0]) &&
SupportMKLDNNAveragePooling(inputs[0], outputs[0])) {
- const AdaptiveAvgPoolParam ¶m =
- nnvm::get<AdaptiveAvgPoolParam>(attrs.parsed);
+ const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed);
+
const NDArray *workspace = nullptr;
MKLDNN_OPCHECK_INIT(false, 1, inputs, outputs);
- MKLDNNAdaptivePoolingCompute(ctx, param, inputs[0], req[0], outputs[0],
- workspace);
+ MKLDNNPoolingCompute(ctx, param, inputs[0], req[0], outputs[0], workspace,
true);
+ MKLDNN_OPCHECK_RUN(PoolingCompute<cpu>, attrs, ctx, inputs, req, outputs);
return;
}
FallBackCompute(AdaptiveAvgPoolOpForward<cpu>, attrs, ctx, inputs, req,
@@ -283,8 +286,6 @@ inline static bool AdaptivePoolingStorageType(const
nnvm::NodeAttrs &attrs,
}
-DMLC_REGISTER_PARAMETER(AdaptiveAvgPoolParam);
-
NNVM_REGISTER_OP(_contrib_AdaptiveAvgPooling2D)
.describe(R"code(
Applies a 2D adaptive average pooling over a 4D input with the shape of (NCHW).
@@ -297,7 +298,7 @@ The pooling kernel and stride sizes are automatically
chosen for desired output
(N x C x height x width) for any input (NCHW).
)code" ADD_FILELINE)
-.set_attr_parser(ParamParser<AdaptiveAvgPoolParam>)
+.set_attr_parser(ParamParser<PoolingParam>)
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr<mxnet::FInferShape>("FInferShape", AdaptiveAvgPoolOpInferShape)
@@ -310,10 +311,10 @@ The pooling kernel and stride sizes are automatically
chosen for desired output
.set_attr<FComputeEx>("FComputeEx<cpu>", AdaptiveAvgPoolComputeExCPU)
#endif
.add_argument("data", "NDArray-or-Symbol", "Input data")
-.add_arguments(AdaptiveAvgPoolParam::__FIELDS__());
+.add_arguments(PoolingParam::__FIELDS__());
NNVM_REGISTER_OP(_backward_contrib_AdaptiveAvgPooling2D)
-.set_attr_parser(ParamParser<AdaptiveAvgPoolParam>)
+.set_attr_parser(ParamParser<PoolingParam>)
.set_num_inputs(1)
.set_num_outputs(1)
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
diff --git a/src/operator/nn/mkldnn/mkldnn_adaptive_pooling-inl.h
b/src/operator/nn/mkldnn/mkldnn_adaptive_pooling-inl.h
deleted file mode 100644
index d46bd59..0000000
--- a/src/operator/nn/mkldnn/mkldnn_adaptive_pooling-inl.h
+++ /dev/null
@@ -1,155 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * Copyright (c) 2021 by Contributors
- * \file mkldnn_adaptive_pooling-inl.h
- */
-#ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_ADAPTIVE_POOLING_INL_H_
-#define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_ADAPTIVE_POOLING_INL_H_
-
-#if MXNET_USE_MKLDNN == 1
-
-#include <mkldnn.hpp>
-#include <utility>
-
-#include "../../operator_common.h"
-#include "./mkldnn_base-inl.h"
-
-namespace mxnet {
-namespace op {
-
-class MKLDNNAdaptivePoolingFwd {
- public:
- MKLDNNAdaptivePoolingFwd(const mxnet::NDArray& input,
- const mxnet::NDArray& output,
- const mkldnn::memory::dims& kernel,
- const mkldnn::memory::dims& strides,
- const mkldnn::memory::dims& pad_l,
- const mkldnn::memory::dims& pad_r,
- const mkldnn::algorithm alg_kind,
- const bool with_workspace,
- const bool is_train)
- : with_workspace_(with_workspace), fwd_(nullptr) {
- Init(input, output, kernel, strides, pad_l, pad_r, is_train, alg_kind);
- }
- ~MKLDNNAdaptivePoolingFwd() = default;
-
- public:
- void Execute(const NDArray& input,
- const OpReqType req,
- const NDArray& output,
- const NDArray* workspace);
-
- private:
- bool with_workspace_;
- std::shared_ptr<mkldnn::pooling_forward::primitive_desc> fwd_pd_;
- std::shared_ptr<mkldnn::pooling_forward> fwd_;
-
- private:
- void Init(const mxnet::NDArray& input,
- const mxnet::NDArray& output,
- const mkldnn::memory::dims& kernel,
- const mkldnn::memory::dims& strides,
- const mkldnn::memory::dims& pad_l,
- const mkldnn::memory::dims& pad_r,
- const bool is_train,
- const mkldnn::algorithm alg_kind);
-};
-
-template <typename T = mkldnn::memory::dims>
-void updateAdaptivePaddingKernel(T* kernel,
- T* strides,
- T* pad_l,
- T* pad_r,
- const NDArray& in_data,
- const NDArray& out_data) {
- const int IH = in_data.shape()[2];
- const int IW = in_data.shape()[3];
- const int OH = out_data.shape()[2];
- const int OW = out_data.shape()[3];
-
- strides->at(0) = floor((IH << 1) / OH) - floor(IH / OH);
- strides->at(1) = floor((IW << 1) / OW) - floor(IW / OW);
- kernel->at(0) = ceil((IH << 1) / OH) - floor(IH / OH);
- kernel->at(1) = ceil((IW << 1) / OW) - floor(IW / OW);
- pad_l->at(0) = (strides->at(0) * (OH - 1) + kernel->at(0) - IH) >> 1;
- pad_l->at(1) = (strides->at(1) * (OW - 1) + kernel->at(1) - IW) >> 1;
-}
-
-template <typename T>
-MKLDNNAdaptivePoolingFwd& GetPoolingFwd(const T& param,
- const bool is_train,
- const NDArray& input,
- const NDArray& output) {
- if (input.shape().ndim() != 4) {
- LOG(FATAL) << "MKLDNN Adaptive Avg Pool 2d: Expect only 2D input";
- }
- typedef ParamOpSign<T> MKLDNNPoolingSignature;
-#if DMLC_CXX11_THREAD_LOCAL
- static thread_local std::unordered_map<MKLDNNPoolingSignature,
MKLDNNAdaptivePoolingFwd, OpHash>
- pooling_fwds;
-#else
- static MX_THREAD_LOCAL
- std::unordered_map<MKLDNNPoolingSignature, MKLDNNAdaptivePoolingFwd,
OpHash>
- pooling_fwds;
-#endif
- bool with_workspace = is_train;
- MKLDNNPoolingSignature key(param);
- key.AddSign(is_train);
- key.AddSign(with_workspace);
- key.AddSign(input);
- key.AddSign(output);
-
- auto it = pooling_fwds.find(key);
- if (it == pooling_fwds.end()) {
- const int kernel_ndims = input.shape().ndim();
-
- mkldnn::memory::dims kernel(kernel_ndims);
- mkldnn::memory::dims strides(kernel_ndims);
- mkldnn::memory::dims pad_l(kernel_ndims);
- mkldnn::memory::dims pad_r(kernel_ndims);
-
- updateAdaptivePaddingKernel(&kernel, &strides, &pad_l, &pad_r, input,
output);
- mkldnn::memory::validate_dims(kernel);
- mkldnn::memory::validate_dims(strides);
- mkldnn::memory::validate_dims(pad_l);
- mkldnn::memory::validate_dims(pad_r);
-
- mkldnn::algorithm kind = mkldnn::algorithm::pooling_avg;
- MKLDNNAdaptivePoolingFwd fwd(input, output, kernel, kernel, pad_l, pad_r,
kind, false, false);
- it = AddToCache(&pooling_fwds, key, fwd);
- }
- return it->second;
-}
-
-template <typename T>
-void MKLDNNAdaptivePoolingCompute(const OpContext& ctx,
- const T& param,
- const NDArray& in_data,
- const OpReqType req,
- const NDArray& out_data,
- const NDArray* workspace) {
- auto& fwd = GetPoolingFwd(param, ctx.is_train, in_data, out_data);
- fwd.Execute(in_data, req, out_data, workspace);
-}
-} // namespace op
-} // namespace mxnet
-#endif // MXNET_USE_MKLDNN == 1
-#endif // MXNET_OPERATOR_NN_MKLDNN_MKLDNN_ADAPTIVE_POOLING_INL_H_
diff --git a/src/operator/nn/mkldnn/mkldnn_adaptive_pooling.cc
b/src/operator/nn/mkldnn/mkldnn_adaptive_pooling.cc
deleted file mode 100644
index 13c842e..0000000
--- a/src/operator/nn/mkldnn/mkldnn_adaptive_pooling.cc
+++ /dev/null
@@ -1,100 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * Copyright (c) 2021 by Contributors
- * \file mkldnn_adaptive_pooling.cc
- */
-
-#if MXNET_USE_MKLDNN == 1
-
-#include "./mkldnn_adaptive_pooling-inl.h"
-
-namespace mxnet {
-namespace op {
-void MKLDNNAdaptivePoolingFwd::Init(const mxnet::NDArray& input,
- const mxnet::NDArray& output,
- const mkldnn::memory::dims& kernel,
- const mkldnn::memory::dims& strides,
- const mkldnn::memory::dims& pad_l,
- const mkldnn::memory::dims& pad_r,
- const bool is_train,
- const mkldnn::algorithm alg_kind) {
- const mkldnn::memory* mem = static_cast<const
mkldnn::memory*>(input.GetMKLDNNData());
- const auto src_md = mem->get_desc();
- const auto dst_md = GetMemDesc(output);
- const mkldnn::engine engine = CpuEngine::Get()->get_engine();
-
- if (alg_kind != mkldnn::algorithm::pooling_avg &&
- alg_kind != mkldnn::algorithm::pooling_avg_include_padding &&
- alg_kind != mkldnn::algorithm::pooling_avg_exclude_padding) {
- LOG(FATAL) << "MKLDNN Adaptive Pooling: algorithm is not supported";
- }
-
- mkldnn::prop_kind prop = mkldnn::prop_kind::forward_scoring;
- if (is_train && alg_kind != mkldnn::algorithm::pooling_avg) {
- prop = mkldnn::prop_kind::forward_training;
- }
- if (is_train && prop == mkldnn::prop_kind::forward_scoring) {
- LOG(INFO) << "MKLDNN Pooling: training with prop_kind is forward_scoring";
- }
-
- const auto fwd_desc =
- mkldnn::pooling_forward::desc(prop, alg_kind, src_md, dst_md, strides,
kernel, pad_l, pad_r);
- this->fwd_pd_.reset(new mkldnn::pooling_forward::primitive_desc(fwd_desc,
engine));
- this->fwd_.reset(new mkldnn::pooling_forward(*(this->fwd_pd_)));
-}
-
-void MKLDNNAdaptivePoolingFwd::Execute(const NDArray& input,
- const OpReqType req,
- const NDArray& output,
- const NDArray* workspace) {
- NDArray in_buffer = input;
- if (input.IsView() && input.IsMKLDNNData()) {
- in_buffer = input.Reorder2Default();
- }
-
- auto input_mem = static_cast<const
mkldnn::memory*>(in_buffer.GetMKLDNNData());
- auto output_mem_t = CreateMKLDNNMem(output, this->fwd_pd_->dst_desc(), req);
-
- mkldnn_args_map_t args = {{MKLDNN_ARG_SRC, *input_mem}, {MKLDNN_ARG_DST,
*(output_mem_t.second)}};
-
- if (this->with_workspace_) {
- auto engine = CpuEngine::Get()->get_engine();
- if (workspace == nullptr) {
- LOG(FATAL) << "MKLDNN Average Pooling: incorrect worskapce input";
- }
- auto ws = std::make_shared<mkldnn::memory>(
- this->fwd_pd_->workspace_desc(),
- engine,
- static_cast<const
mkldnn::memory*>(workspace->GetMKLDNNData())->get_data_handle());
- args[MKLDNN_ARG_WORKSPACE] = *ws;
- }
- if (this->fwd_) {
- MKLDNNStream::Get()->RegisterPrimArgs(*(this->fwd_), args);
- CommitOutput(output, output_mem_t);
- MKLDNNStream::Get()->Submit();
- } else {
- LOG(FATAL) << "MKLDNN Pooling: forward primitive is nullptr";
- }
-}
-
-} // namespace op
-} // namespace mxnet
-#endif // MXNET_USE_MKLDNN == 1
diff --git a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
index bf9cf46..775c6a9 100644
--- a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
+++ b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
@@ -87,6 +87,26 @@ class MKLDNNPoolingBwd {
const mkldnn::pooling_backward::primitive_desc& GetPd();
};
+template <typename T = mkldnn::memory::dims>
+void UseAdaptivePaddingKernel(T* kernel,
+ T* strides,
+ T* pad_l,
+ T* pad_r,
+ const NDArray& in_data,
+ const NDArray& out_data) {
+ const int IH = in_data.shape()[2];
+ const int IW = in_data.shape()[3];
+ const int OH = out_data.shape()[2];
+ const int OW = out_data.shape()[3];
+
+ strides->at(0) = floor((IH << 1) / OH) - floor(IH / OH);
+ strides->at(1) = floor((IW << 1) / OW) - floor(IW / OW);
+ kernel->at(0) = ceil((IH << 1) / OH) - floor(IH / OH);
+ kernel->at(1) = ceil((IW << 1) / OW) - floor(IW / OW);
+ pad_l->at(0) = (strides->at(0) * (OH - 1) + kernel->at(0) - IH) >> 1;
+ pad_l->at(1) = (strides->at(1) * (OW - 1) + kernel->at(1) - IW) >> 1;
+}
+
inline int GetPaddingSizeFull(dim_t x, int padl, int padr, int k, int s) {
if ((x + padl + padr - k) % s != 0) {
return (padr + s - ((x + padl + padr - k) % s));
@@ -158,7 +178,8 @@ void MKLDNNPoolingCompute(const OpContext& ctx,
const NDArray& in_data,
const OpReqType req,
const NDArray& out_data,
- const NDArray* workspace);
+ const NDArray* workspace,
+ const bool use_adaptive_pooling);
void MKLDNNPoolingGradCompute(const OpContext& ctx,
const PoolingParam& param,
@@ -170,7 +191,8 @@ void MKLDNNPoolingGradCompute(const OpContext& ctx,
MKLDNNPoolingFwd& GetPoolingFwd(const PoolingParam& param,
const bool is_train,
const NDArray& data,
- const NDArray& output);
+ const NDArray& output,
+ const bool use_adaptive_pooling);
} // namespace op
} // namespace mxnet
#endif // MXNET_USE_MKLDNN == 1
diff --git a/src/operator/nn/mkldnn/mkldnn_pooling.cc
b/src/operator/nn/mkldnn/mkldnn_pooling.cc
index 9a88b86..d9500bc 100644
--- a/src/operator/nn/mkldnn/mkldnn_pooling.cc
+++ b/src/operator/nn/mkldnn/mkldnn_pooling.cc
@@ -45,6 +45,7 @@ void MKLDNNPoolingFwd::Init(const mxnet::NDArray& input,
const auto src_md = static_cast<const
mkldnn::memory*>(input.GetMKLDNNData())->get_desc();
const auto dst_md = GetMemDesc(output);
const mkldnn::engine engine = CpuEngine::Get()->get_engine();
+
if (alg_kind != mkldnn::algorithm::pooling_max && alg_kind !=
mkldnn::algorithm::pooling_avg &&
alg_kind != mkldnn::algorithm::pooling_avg_include_padding &&
alg_kind != mkldnn::algorithm::pooling_avg_exclude_padding) {
@@ -107,7 +108,7 @@ void MKLDNNPoolingFwd::Execute(const NDArray& in_data,
}
}
-mkldnn::algorithm GetMKLDNNPoolAlgo(const PoolingParam& param) {
+mkldnn::algorithm GetMKLDNNPoolingAlgorithm(const PoolingParam& param) {
switch (param.pool_type) {
case pool_enum::kMaxPooling:
return mkldnn::algorithm::pooling_max;
@@ -201,7 +202,7 @@ mkldnn::pooling_forward::primitive_desc
GetPoolingFwdPdesc(const PoolingParam& p
InitPoolingPrimitiveParams(param, data_md, kernel, strides, pad_l, pad_r);
- const mkldnn::algorithm alg = GetMKLDNNPoolAlgo(param);
+ const mkldnn::algorithm alg = GetMKLDNNPoolingAlgorithm(param);
mkldnn::prop_kind kind = mkldnn::prop_kind::forward_scoring;
if (is_train && alg != mkldnn::algorithm::pooling_avg) {
kind = mkldnn::prop_kind::forward_training;
@@ -215,7 +216,8 @@ mkldnn::pooling_forward::primitive_desc
GetPoolingFwdPdesc(const PoolingParam& p
MKLDNNPoolingFwd& GetPoolingFwd(const PoolingParam& param,
const bool is_train,
const NDArray& data,
- const NDArray& output) {
+ const NDArray& output,
+ const bool use_adaptive_pooling) {
#if DMLC_CXX11_THREAD_LOCAL
static thread_local std::unordered_map<MKLDNNPoolingSignature,
MKLDNNPoolingFwd, OpHash>
pooling_fwds;
@@ -224,27 +226,41 @@ MKLDNNPoolingFwd& GetPoolingFwd(const PoolingParam& param,
pooling_fwds;
#endif
- bool with_workspace = is_train && MKLDNNRequireWorkspace(param);
+ const bool with_workspace = is_train && MKLDNNRequireWorkspace(param);
MKLDNNPoolingSignature key(param);
key.AddSign(is_train);
key.AddSign(with_workspace);
key.AddSign(data);
key.AddSign(output);
-
+ if (use_adaptive_pooling) {
+ key.AddSign(use_adaptive_pooling);
+ }
auto it = pooling_fwds.find(key);
if (it == pooling_fwds.end()) {
- CHECK(param.kernel.ndim() == 1 || param.kernel.ndim() == 2 ||
param.kernel.ndim() == 3)
+ CHECK(use_adaptive_pooling || (param.kernel.ndim() >= 1 &&
param.kernel.ndim() <= 3))
<< "Not Implemented";
auto data_md = static_cast<const
mkldnn::memory*>(data.GetMKLDNNData())->get_desc();
- const auto kernel_ndims = param.kernel.ndim();
+ const auto kernel_ndims = use_adaptive_pooling ? data.shape().ndim() :
param.kernel.ndim();
+
mkldnn::memory::dims kernel(kernel_ndims);
mkldnn::memory::dims strides(kernel_ndims);
mkldnn::memory::dims pad_l(kernel_ndims);
mkldnn::memory::dims pad_r(kernel_ndims);
- InitPoolingPrimitiveParams(param, data_md, kernel, strides, pad_l, pad_r);
- const mkldnn::algorithm alg = GetMKLDNNPoolAlgo(param);
+ if (use_adaptive_pooling) {
+ UseAdaptivePaddingKernel(&kernel, &strides, &pad_l, &pad_r, data,
output);
+ mkldnn::memory::validate_dims(kernel);
+ mkldnn::memory::validate_dims(strides);
+ mkldnn::memory::validate_dims(pad_l);
+ mkldnn::memory::validate_dims(pad_r);
+ } else {
+ InitPoolingPrimitiveParams(param, data_md, kernel, strides, pad_l,
pad_r);
+ }
+
+ const mkldnn::algorithm alg =
+ use_adaptive_pooling ? mkldnn::algorithm::pooling_avg :
GetMKLDNNPoolingAlgorithm(param);
+
MKLDNNPoolingFwd fwd(
data, output, kernel, strides, pad_l, pad_r, alg, with_workspace,
is_train);
it = AddToCache(&pooling_fwds, key, fwd);
@@ -257,8 +273,9 @@ void MKLDNNPoolingCompute(const OpContext& ctx,
const NDArray& in_data,
const OpReqType req,
const NDArray& out_data,
- const NDArray* workspace) {
- auto& fwd = GetPoolingFwd(param, ctx.is_train, in_data, out_data);
+ const NDArray* workspace,
+ const bool use_adaptive_pooling) {
+ auto& fwd = GetPoolingFwd(param, ctx.is_train, in_data, out_data,
use_adaptive_pooling);
fwd.Execute(in_data, req, out_data, workspace);
}
@@ -306,7 +323,7 @@ MKLDNNPoolingBwd& GetPoolingBwd(const PoolingParam& param,
auto diff_src_dims = mkldnn::memory::dims(in_grad.shape().begin(),
in_grad.shape().end());
auto diff_src_md = mkldnn::memory::desc(diff_src_dims,
get_data_type(data_md), any);
auto cpu_engine = CpuEngine::Get()->get_engine();
- auto alg = GetMKLDNNPoolAlgo(param);
+ auto alg = GetMKLDNNPoolingAlgorithm(param);
const int kernel_ndims = param.kernel.ndim();
mkldnn::memory::dims kernel(kernel_ndims);
diff --git a/src/operator/nn/pooling-inl.h b/src/operator/nn/pooling-inl.h
index 03f0fa8..1c83ec5 100644
--- a/src/operator/nn/pooling-inl.h
+++ b/src/operator/nn/pooling-inl.h
@@ -54,6 +54,8 @@ struct PoolingParam : public dmlc::Parameter<PoolingParam> {
dmlc::optional<int> p_value;
dmlc::optional<bool> count_include_pad;
dmlc::optional<int> layout;
+ dmlc::optional<mxnet::Tuple<int>> output_size;
+
DMLC_DECLARE_PARAMETER(PoolingParam) {
DMLC_DECLARE_FIELD(kernel).set_default(mxnet::TShape(0, 0)) // add
default value here
.enforce_nonzero()
@@ -104,6 +106,10 @@ struct PoolingParam : public dmlc::Parameter<PoolingParam>
{
.set_default(dmlc::optional<int>())
.describe("Set layout for input and output. Empty for\n "
"default layout: NCW for 1d, NCHW for 2d and NCDHW for 3d.");
+
+
DMLC_DECLARE_FIELD(output_size).set_default(dmlc::optional<mxnet::Tuple<int>>())
+ .describe("Only used for Adaptive Pooling. int (output size) or a tuple of
int for output "
+ "(height, width).");
}
bool operator==(const PoolingParam& other) const {
@@ -116,7 +122,8 @@ struct PoolingParam : public dmlc::Parameter<PoolingParam> {
this->cudnn_off == other.cudnn_off &&
this->p_value == other.p_value &&
this->count_include_pad == other.count_include_pad &&
- this->layout == other.layout;
+ this->layout == other.layout &&
+ this->output_size == other.output_size;
}
// Extract layout from param, or supply default layout based on provided
input dimension.
diff --git a/src/operator/nn/pooling.cc b/src/operator/nn/pooling.cc
index c81cae3..4a6d105 100644
--- a/src/operator/nn/pooling.cc
+++ b/src/operator/nn/pooling.cc
@@ -268,7 +268,8 @@ static bool PoolingShape(const nnvm::NodeAttrs &attrs,
}
#if MXNET_USE_MKLDNN == 1
-void PoolingComputeExCPU(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
+void PoolingComputeExCPU(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
const std::vector<NDArray> &inputs,
const std::vector<OpReqType> &req,
const std::vector<NDArray> &outputs) {
@@ -287,7 +288,7 @@ void PoolingComputeExCPU(const nnvm::NodeAttrs &attrs,
const OpContext &ctx,
workspace = &outputs[1];
}
MKLDNN_OPCHECK_INIT(false, 1, inputs, outputs);
- MKLDNNPoolingCompute(ctx, param, inputs[0], req[0], outputs[0], workspace);
+ MKLDNNPoolingCompute(ctx, param, inputs[0], req[0], outputs[0], workspace,
false);
MKLDNN_OPCHECK_RUN(PoolingCompute<cpu>, attrs, ctx, inputs, req, outputs);
return;
}
diff --git a/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
b/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
index e7d50ef..98d1754 100644
--- a/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
+++ b/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
@@ -39,7 +39,8 @@ static void MKLDNNQuantizedPoolingForward(const
nnvm::NodeAttrs& attrs,
<< "mkldnn_quantized_pooling op only supports uint8 and int8 as input "
"type";
const PoolingParam& param = nnvm::get<PoolingParam>(attrs.parsed);
- MKLDNNPoolingCompute(ctx, param, in_data[0], req[0], out_data[0], nullptr);
+ MKLDNNPoolingCompute(
+ ctx, param, in_data[0], req[0], out_data[0], nullptr, /*use_adaptive*/
false);
out_data[1].data().dptr<float>()[0] = in_data[1].data().dptr<float>()[0];
out_data[2].data().dptr<float>()[0] = in_data[2].data().dptr<float>()[0];
}
