| /* |
| * 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) 2017 by Contributors |
| * \file pooling.cc |
| * \brief |
| * \author Bing Xu, Jun Wu, Da Zheng |
| */ |
| #include "../elemwise_op_common.h" |
| #include "./pooling-inl.h" |
| #if MXNET_USE_NNPACK == 1 |
| #include "./nnpack/nnpack_pooling-inl.h" |
| #endif // MXNET_USE_NNPACK |
| #if MXNET_USE_MKLDNN == 1 |
| #include "./mkldnn/mkldnn_pooling-inl.h" |
| #endif // MXNET_USE_MKLDNN |
| |
| namespace mxnet { |
| namespace op { |
| |
| static void PoolingParamParser(nnvm::NodeAttrs *attrs) { |
| using namespace mshadow; |
| PoolingParam param; |
| param.Init(attrs->dict); |
| if (param.kernel.ndim() == 1) { |
| if (param.stride.ndim() == 0) param.stride = Shape1(1); |
| if (param.pad.ndim() == 0) param.pad = Shape1(0); |
| } else if (param.kernel.ndim() == 2) { |
| if (param.stride.ndim() == 0) param.stride = Shape2(1, 1); |
| if (param.pad.ndim() == 0) param.pad = Shape2(0, 0); |
| } else { |
| CHECK_EQ(param.kernel.ndim(), 3U) << param.kernel.ndim() |
| << "D pooling not supported"; |
| if (param.stride.ndim() == 0) param.stride = Shape3(1, 1, 1); |
| if (param.pad.ndim() == 0) param.pad = Shape3(0, 0, 0); |
| } |
| CHECK_EQ(param.stride.ndim(), param.kernel.ndim()) |
| << "stride and kernel should have the same length"; |
| CHECK_EQ(param.pad.ndim(), param.kernel.ndim()) |
| << "pad and kernel should have the same length"; |
| attrs->parsed = std::move(param); |
| } |
| |
| int GetNumOutputs(const PoolingParam ¶m) { |
| #if MXNET_USE_MKLDNN == 1 |
| return MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param) ? 2 : 1; |
| #else |
| return 1; |
| #endif |
| } |
| |
| int GetNumBackInputs(const PoolingParam ¶m) { |
| #if MXNET_USE_MKLDNN == 1 |
| return MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param) ? 5 : 3; |
| #else |
| return 3; |
| #endif |
| } |
| |
| static bool PoolingType(const nnvm::NodeAttrs& attrs, |
| std::vector<int> *in_attrs, |
| std::vector<int> *out_attrs) { |
| out_attrs->at(0) = in_attrs->at(0); |
| #if MXNET_USE_MKLDNN == 1 |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| if (MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param)) { |
| CHECK_GT(out_attrs->size(), 1U); |
| out_attrs->at(1) = mshadow::kInt32; |
| } |
| #endif |
| return true; |
| } |
| |
| static bool PoolingShape(const nnvm::NodeAttrs &attrs, |
| std::vector<TShape> *in_shape, |
| std::vector<TShape> *out_shape) { |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| CHECK_EQ(in_shape->size(), 1U); |
| const TShape &dshape = (*in_shape)[0]; |
| CHECK_GE(dshape.ndim(), 3U) |
| << "Pooling: Input data should be 3D in (batch, channel, x)" |
| << " Or 4D in (batch, channel, y, x) " |
| << " Or 5D in (batch, channel, d, y, x)"; |
| TShape oshape = dshape; |
| if (dshape.ndim() == 0) return false; |
| if (param.kernel.ndim() == 1) { |
| CHECK_EQ(dshape.ndim(), 3U) |
| << "Pooling: Input data should be 3D in (batch, channel, x)"; |
| if (param.global_pool) { |
| oshape[2] = 1; |
| } else { |
| CHECK(param.kernel[0] <= dshape[2] + 2 * param.pad[0]) |
| << "kernel size (" << param.kernel[0] << ") exceeds input (" |
| << dshape[2] << " padded to " << (dshape[2] + 2 * param.pad[0]) |
| << ")"; |
| if (param.pooling_convention == pool_enum::kValid) { |
| oshape[2] = 1 + |
| (dshape[2] + 2 * param.pad[0] - param.kernel[0]) / |
| param.stride[0]; |
| } else { |
| oshape[2] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[2] + 2 * param.pad[0] - |
| param.kernel[0]) / |
| param.stride[0])); |
| } |
| } |
| out_shape->clear(); |
| out_shape->push_back(oshape); // save output shape |
| #if MXNET_USE_MKLDNN == 1 |
| if (MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param)) |
| out_shape->push_back(oshape); // for workspace |
| #endif |
| } else if (param.kernel.ndim() == 2) { |
| CHECK_EQ(dshape.ndim(), 4U) |
| << "Pooling: Input data should be 4D in (batch, channel, y, x)"; |
| if (param.global_pool) { |
| oshape[2] = 1; |
| oshape[3] = 1; |
| } else { |
| CHECK(param.kernel[0] <= dshape[2] + 2 * param.pad[0]) |
| << "kernel size (" << param.kernel[0] << ") exceeds input (" |
| << dshape[2] << " padded to " << (dshape[2] + 2 * param.pad[0]) |
| << ")"; |
| CHECK(param.kernel[1] <= dshape[3] + 2 * param.pad[1]) |
| << "kernel size (" << param.kernel[1] << ") exceeds input (" |
| << dshape[3] << " padded to " << (dshape[3] + 2 * param.pad[1]) |
| << ")"; |
| if (param.pooling_convention == pool_enum::kValid) { |
| oshape[2] = 1 + |
| (dshape[2] + 2 * param.pad[0] - param.kernel[0]) / |
| param.stride[0]; |
| oshape[3] = 1 + |
| (dshape[3] + 2 * param.pad[1] - param.kernel[1]) / |
| param.stride[1]; |
| } else { |
| oshape[2] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[2] + 2 * param.pad[0] - |
| param.kernel[0]) / |
| param.stride[0])); |
| oshape[3] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[3] + 2 * param.pad[1] - |
| param.kernel[1]) / |
| param.stride[1])); |
| } |
| } |
| out_shape->clear(); |
| out_shape->push_back(oshape); // save output shape |
| #if MXNET_USE_MKLDNN == 1 |
| if (MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param)) |
| out_shape->push_back(oshape); // for workspace |
| #endif |
| } else if (param.kernel.ndim() == 3) { |
| CHECK_EQ(dshape.ndim(), 5U) |
| << "Pooling: Input data should be 5D in (batch, channel, d, y, x)"; |
| CHECK_LE(param.kernel[0], dshape[2] + 2 * param.pad[0]) |
| << "kernel size exceeds input"; |
| CHECK_LE(param.kernel[1], dshape[3] + 2 * param.pad[1]) |
| << "kernel size exceeds input"; |
| CHECK_LE(param.kernel[2], dshape[4] + 2 * param.pad[2]) |
| << "kernel size exceeds input"; |
| if (param.global_pool) { |
| oshape[2] = 1; |
| oshape[3] = 1; |
| oshape[4] = 1; |
| } else { |
| if (param.pooling_convention == pool_enum::kValid) { |
| oshape[2] = 1 + |
| (dshape[2] + 2 * param.pad[0] - param.kernel[0]) / |
| param.stride[0]; |
| oshape[3] = 1 + |
| (dshape[3] + 2 * param.pad[1] - param.kernel[1]) / |
| param.stride[1]; |
| oshape[4] = 1 + |
| (dshape[4] + 2 * param.pad[2] - param.kernel[2]) / |
| param.stride[2]; |
| } else { |
| oshape[2] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[2] + 2 * param.pad[0] - |
| param.kernel[0]) / |
| param.stride[0])); |
| oshape[3] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[3] + 2 * param.pad[1] - |
| param.kernel[1]) / |
| param.stride[1])); |
| oshape[4] = 1 + static_cast<int>(ceil( |
| static_cast<float>(dshape[4] + 2 * param.pad[2] - |
| param.kernel[2]) / |
| param.stride[2])); |
| } |
| } |
| |
| out_shape->clear(); |
| out_shape->push_back(oshape); // save output shape |
| #if MXNET_USE_MKLDNN == 1 |
| if (MKLDNNRequireWorkspace(param) && SupportMKLDNNPooling(param)) |
| out_shape->push_back(oshape); // for workspace |
| #endif |
| } |
| return true; |
| } |
| |
| #if MXNET_USE_MKLDNN == 1 |
| void PoolingComputeExCPU(const nnvm::NodeAttrs &attrs, const OpContext &ctx, |
| const std::vector<NDArray> &inputs, |
| const std::vector<OpReqType> &req, |
| const std::vector<NDArray> &outputs) { |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| const NDArray *workspace = nullptr; |
| if (MKLDNNRequireWorkspace(param)) { |
| CHECK_GT(outputs.size(), 1U); |
| workspace = &outputs[1]; |
| } |
| if (SupportMKLDNN(inputs[0]) |
| && SupportMKLDNNPooling(param, inputs[0].shape())) { |
| MKLDNN_OPCHECK_INIT(false, 1, inputs, outputs); |
| MKLDNNPoolingCompute(ctx, param, inputs[0], req[0], outputs[0], workspace); |
| MKLDNN_OPCHECK_RUN(PoolingCompute<cpu>, attrs, ctx, inputs, req, outputs); |
| return; |
| } |
| FallBackCompute(PoolingCompute<cpu>, attrs, ctx, inputs, req, outputs); |
| } |
| |
| void PoolingGradComputeExCPU(const nnvm::NodeAttrs &attrs, const OpContext &ctx, |
| const std::vector<NDArray> &inputs, |
| const std::vector<OpReqType> &req, |
| const std::vector<NDArray> &outputs) { |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| const NDArray &out_grad = inputs[0]; |
| const NDArray *workspace = nullptr; |
| const NDArray *in_data = nullptr; |
| if (MKLDNNRequireWorkspace(param)) { |
| // The first two elements are the gradient of the outputs in forward. |
| // The third is the input of forward. |
| // The fourth and the fifth are the outputs of forward. |
| CHECK_EQ(inputs.size(), 5U); |
| in_data = &inputs[2]; |
| workspace = &inputs[4]; |
| } else { |
| CHECK_EQ(inputs.size(), 3U); |
| in_data = &inputs[1]; |
| } |
| const NDArray &in_grad = outputs[0]; |
| if (SupportMKLDNN(inputs[0]) |
| && SupportMKLDNNPooling(param, inputs[0].shape())) { |
| MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs); |
| MKLDNNPoolingGradCompute(ctx, param, out_grad, *in_data, workspace, |
| req[0], in_grad); |
| MKLDNN_OPCHECK_RUN(PoolingGradCompute<cpu>, attrs, ctx, inputs, req, |
| outputs); |
| return; |
| } |
| FallBackCompute(PoolingGradCompute<cpu>, attrs, ctx, inputs, req, outputs); |
| } |
| #endif |
| |
| inline static bool PoolingStorageType(const nnvm::NodeAttrs &attrs, |
| const int dev_mask, |
| DispatchMode *dispatch_mode, |
| std::vector<int> *in_attrs, |
| std::vector<int> *out_attrs) { |
| CHECK_EQ(in_attrs->size(), 1); |
| |
| #if MXNET_USE_MKLDNN == 1 |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNPooling(param)) { |
| return storage_type_assign(out_attrs, mxnet::kDefaultStorage, |
| dispatch_mode, DispatchMode::kFComputeEx); |
| } |
| #else |
| CHECK_EQ(out_attrs->size(), 1); |
| #endif |
| return storage_type_assign(out_attrs, mxnet::kDefaultStorage, |
| dispatch_mode, DispatchMode::kFCompute); |
| } |
| |
| inline static bool BackwardPoolingStorageType(const nnvm::NodeAttrs &attrs, |
| const int dev_mask, |
| DispatchMode *dispatch_mode, |
| std::vector<int> *in_attrs, |
| std::vector<int> *out_attrs) { |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| CHECK_EQ(in_attrs->size(), GetNumBackInputs(param)); |
| CHECK_EQ(out_attrs->size(), 1); |
| |
| #if MXNET_USE_MKLDNN == 1 |
| if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNPooling(param)) { |
| return storage_type_assign(out_attrs, mxnet::kDefaultStorage, |
| dispatch_mode, DispatchMode::kFComputeEx); |
| } |
| #else |
| CHECK_EQ(in_attrs->size(), 3); |
| #endif |
| return storage_type_assign(out_attrs, mxnet::kDefaultStorage, |
| dispatch_mode, DispatchMode::kFCompute); |
| } |
| |
| DMLC_REGISTER_PARAMETER(PoolingParam); |
| |
| NNVM_REGISTER_OP(Pooling) |
| .describe(R"code(Performs pooling on the input. |
| |
| The shapes for 1-D pooling are |
| |
| - **data**: *(batch_size, channel, width)*, |
| - **out**: *(batch_size, num_filter, out_width)*. |
| |
| The shapes for 2-D pooling are |
| |
| - **data**: *(batch_size, channel, height, width)* |
| - **out**: *(batch_size, num_filter, out_height, out_width)*, with:: |
| |
| out_height = f(height, kernel[0], pad[0], stride[0]) |
| out_width = f(width, kernel[1], pad[1], stride[1]) |
| |
| The definition of *f* depends on ``pooling_convention``, which has two options: |
| |
| - **valid** (default):: |
| |
| f(x, k, p, s) = floor((x+2*p-k)/s)+1 |
| |
| - **full**, which is compatible with Caffe:: |
| |
| f(x, k, p, s) = ceil((x+2*p-k)/s)+1 |
| |
| But ``global_pool`` is set to be true, then do a global pooling, namely reset |
| ``kernel=(height, width)``. |
| |
| Three pooling options are supported by ``pool_type``: |
| |
| - **avg**: average pooling |
| - **max**: max pooling |
| - **sum**: sum pooling |
| |
| For 3-D pooling, an additional *depth* dimension is added before |
| *height*. Namely the input data will have shape *(batch_size, channel, depth, |
| height, width)*. |
| |
| )code" ADD_FILELINE) |
| .set_num_inputs(1) |
| .set_num_outputs([](const NodeAttrs& attrs) { |
| const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed); |
| return GetNumOutputs(param); |
| }) |
| #if MXNET_USE_MKLDNN == 1 |
| .set_attr<nnvm::FNumVisibleOutputs>("FNumVisibleOutputs", |
| [](const NodeAttrs& attrs) { return 1; }) |
| #endif |
| .set_attr<nnvm::FListInputNames>("FListInputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"data"}; |
| }) |
| .set_attr<nnvm::FListOutputNames>("FListOutputNames", |
| [](const NodeAttrs& attrs) { |
| return std::vector<std::string>{"output"}; |
| }) |
| .set_attr_parser(PoolingParamParser) |
| .set_attr<FInferStorageType>("FInferStorageType", PoolingStorageType) |
| .set_attr<nnvm::FInferType>("FInferType", PoolingType) |
| .set_attr<nnvm::FInferShape>("FInferShape", PoolingShape) |
| .set_attr<FCompute>("FCompute<cpu>", PoolingCompute<cpu>) |
| #if MXNET_USE_MKLDNN == 1 |
| .set_attr<FComputeEx>("FComputeEx<cpu>", PoolingComputeExCPU) |
| #endif |
| .set_attr<nnvm::FGradient>("FGradient", |
| ElemwiseGradUseInOut{"_backward_Pooling"}) |
| .add_argument("data", "NDArray-or-Symbol", |
| "Input data to the pooling operator.") |
| .add_arguments(PoolingParam::__FIELDS__()); |
| |
| NNVM_REGISTER_OP(_backward_Pooling) |
| .set_num_outputs(1) |
| .set_attr<nnvm::TIsBackward>("TIsBackward", true) |
| .set_attr<nnvm::FInplaceOption>( |
| "FInplaceOption", |
| [](const NodeAttrs &attrs) { |
| #if MXNET_USE_CUDNN == 1 |
| return std::vector<std::pair<int, int> >(); |
| #else |
| return std::vector<std::pair<int, int> >{{1, 0}}; |
| #endif |
| }) |
| #if MXNET_USE_MKLDNN == 1 |
| .set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& n) { |
| return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; |
| }) |
| #endif |
| .set_attr<FInferStorageType>("FInferStorageType", |
| BackwardPoolingStorageType) |
| .set_attr_parser(PoolingParamParser) |
| #if MXNET_USE_MKLDNN == 1 |
| .set_attr<FComputeEx>("FComputeEx<cpu>", PoolingGradComputeExCPU) |
| #endif |
| .set_attr<FCompute>("FCompute<cpu>", PoolingGradCompute<cpu>); |
| |
| } // namespace op |
| } // namespace mxnet |
