diff --git a/nnvm/include/nnvm/top/nn.h b/nnvm/include/nnvm/top/nn.h index 80d8e02882ef..ec21abde9b8c 100644 --- a/nnvm/include/nnvm/top/nn.h +++ b/nnvm/include/nnvm/top/nn.h @@ -218,14 +218,14 @@ struct Conv2DTransposeParam : public dmlc::Parameter { }; -struct Pool2DParam : public dmlc::Parameter { +struct MaxPool2DParam : public dmlc::Parameter { TShape pool_size; TShape strides; TShape padding; std::string layout; bool ceil_mode; - DMLC_DECLARE_PARAMETER(Pool2DParam) { + DMLC_DECLARE_PARAMETER(MaxPool2DParam) { DMLC_DECLARE_FIELD(pool_size) .describe("Size of the pooling windows.."); DMLC_DECLARE_FIELD(strides).set_default(TShape({1, 1})) @@ -244,6 +244,35 @@ struct Pool2DParam : public dmlc::Parameter { }; +struct AvgPool2DParam : public dmlc::Parameter { + TShape pool_size; + TShape strides; + TShape padding; + std::string layout; + bool ceil_mode; + bool count_include_pad; + + DMLC_DECLARE_PARAMETER(AvgPool2DParam) { + DMLC_DECLARE_FIELD(pool_size) + .describe("Size of the pooling windows.."); + DMLC_DECLARE_FIELD(strides).set_default(TShape({1, 1})) + .describe("Specifies the strides of the convolution."); + DMLC_DECLARE_FIELD(padding).set_default(TShape({0, 0})) + .describe("If padding is non-zero, then the input is implicitly zero-padded" + "on both sides for padding number of points"); + DMLC_DECLARE_FIELD(layout).set_default("NCHW") + .describe("Dimension ordering of data and weight. Can be 'NCHW', 'NHWC', etc." + "'N', 'C', 'H', 'W' stands for batch, channel, height, and width" + "dimensions respectively. Convolution is applied on the 'H' and" + "'W' dimensions."); + DMLC_DECLARE_FIELD(ceil_mode).set_default(false) + .describe("When true, will use ceil instead of floor to compute the output shape."); + DMLC_DECLARE_FIELD(count_include_pad).set_default(false) + .describe("When true, will include padding to compute the average"); + } +}; + + struct GlobalPool2DParam : public dmlc::Parameter { std::string layout; diff --git a/nnvm/src/compiler/fold_scale_axis.cc b/nnvm/src/compiler/fold_scale_axis.cc index 28796a2b0bcd..e38082b69916 100644 --- a/nnvm/src/compiler/fold_scale_axis.cc +++ b/nnvm/src/compiler/fold_scale_axis.cc @@ -354,28 +354,28 @@ NNVM_REGISTER_OP(leaky_relu) .set_attr("FScaleAxisForward", ReluScaleAxisForward); // property registration. +template bool Pool2DBackward( const NodeAttrs& attrs, const std::vector& in_shape, const std::vector& out_shape, const FoldChainInfo& out_info, std::vector* in_axis) { - using top::Pool2DParam; - const Pool2DParam& param = nnvm::get(attrs.parsed); + const T& param = nnvm::get(attrs.parsed); if (out_info.axis == 1 && param.layout == "NCHW") { (*in_axis)[0] = out_info; } return false; } +template bool Pool2DForward( const NodeAttrs& attrs, const std::vector& in_shape, const std::vector& out_shape, std::vector* in_info, FoldChainInfo* out_info) { - using top::Pool2DParam; - const Pool2DParam& param = nnvm::get(attrs.parsed); + const T& param = nnvm::get(attrs.parsed); if ((*in_info)[0].axis == 1 && param.layout == "NCHW") { *out_info = (*in_info)[0]; } @@ -383,16 +383,16 @@ bool Pool2DForward( } NNVM_REGISTER_OP(max_pool2d) -.set_attr("FScaleAxisBackward", Pool2DBackward); +.set_attr("FScaleAxisBackward", Pool2DBackward); NNVM_REGISTER_OP(avg_pool2d) -.set_attr("FScaleAxisBackward", Pool2DBackward); +.set_attr("FScaleAxisBackward", Pool2DBackward); NNVM_REGISTER_OP(max_pool2d) -.set_attr("FScaleAxisForward", Pool2DForward); +.set_attr("FScaleAxisForward", Pool2DForward); NNVM_REGISTER_OP(avg_pool2d) -.set_attr("FScaleAxisForward", Pool2DForward); +.set_attr("FScaleAxisForward", Pool2DForward); diff --git a/nnvm/src/top/nn/pooling.cc b/nnvm/src/top/nn/pooling.cc index 0c40da520544..54eb0d7db3be 100644 --- a/nnvm/src/top/nn/pooling.cc +++ b/nnvm/src/top/nn/pooling.cc @@ -19,12 +19,13 @@ namespace top { using namespace tvm; using namespace nnvm::compiler; -DMLC_REGISTER_PARAMETER(Pool2DParam); +DMLC_REGISTER_PARAMETER(MaxPool2DParam); +template inline bool Pool2DInferShape(const nnvm::NodeAttrs& attrs, std::vector* in_shape, std::vector* out_shape) { - const Pool2DParam& param = nnvm::get(attrs.parsed); + const T& param = nnvm::get(attrs.parsed); CHECK_EQ(in_shape->size(), 1U); CHECK_EQ(out_shape->size(), 1U); @@ -66,11 +67,12 @@ inline bool Pool2DInferShape(const nnvm::NodeAttrs& attrs, return true; } +template inline bool Pool2DCorrectLayout(const NodeAttrs& attrs, std::vector *ilayouts, const std::vector *last_ilayouts, std::vector *olayouts) { - const Pool2DParam ¶m = nnvm::get(attrs.parsed); + const T ¶m = nnvm::get(attrs.parsed); CHECK_EQ(ilayouts->size(), 1); CHECK_EQ(last_ilayouts->size(), 1); CHECK_EQ(olayouts->size(), 1); @@ -114,18 +116,18 @@ NNVM_REGISTER_OP(max_pool2d) )code" NNVM_ADD_FILELINE) .add_argument("data", "4D Tensor", "Input data.") -.add_arguments(Pool2DParam::__FIELDS__()) -.set_attr_parser(ParamParser) -.set_attr("FGetAttrDict", ParamGetAttrDict) +.add_arguments(MaxPool2DParam::__FIELDS__()) +.set_attr_parser(ParamParser) +.set_attr("FGetAttrDict", ParamGetAttrDict) .set_num_outputs(1) .set_num_inputs(1) -.set_attr("FInferShape", Pool2DInferShape) +.set_attr("FInferShape", Pool2DInferShape) .set_attr("FInferType", ElemwiseType<1, 1>) -.set_attr("FCorrectLayout", Pool2DCorrectLayout) +.set_attr("FCorrectLayout", Pool2DCorrectLayout) .set_attr("FTVMCompute", [](const NodeAttrs& attrs, const Array& inputs, const Array& out_info) { - const Pool2DParam& param = nnvm::get(attrs.parsed); + const MaxPool2DParam& param = nnvm::get(attrs.parsed); auto pool_size = ShapeToArray(param.pool_size); auto strides = ShapeToArray(param.strides); auto padding = ShapeToArray(param.padding); @@ -163,12 +165,13 @@ NNVM_REGISTER_OP(_max_pool2d_grad) .add_argument("output", "4D Tensor", "Output data of max_pool2d grad.") .set_num_inputs(3) .set_num_outputs(1) -.set_attr_parser(ParamParser) -.set_attr("FGetAttrDict", ParamGetAttrDict) +.set_attr_parser(ParamParser) +.set_attr("FGetAttrDict", ParamGetAttrDict) .set_attr("FInferShape", AssignOutputAttr) .set_attr("FInferType", ElemwiseType<3, 1>) .set_attr("TIsBackward", true); +DMLC_REGISTER_PARAMETER(AvgPool2DParam); NNVM_REGISTER_OP(avg_pool2d) .describe(R"code(Average pooling operation for one dimensional data. @@ -187,20 +190,21 @@ NNVM_REGISTER_OP(avg_pool2d) )code" NNVM_ADD_FILELINE) .add_argument("data", "4D Tensor", "Input data.") -.add_arguments(Pool2DParam::__FIELDS__()) -.set_attr_parser(ParamParser) -.set_attr("FGetAttrDict", ParamGetAttrDict) -.set_attr("FInferShape", Pool2DInferShape) +.add_arguments(AvgPool2DParam::__FIELDS__()) +.set_attr_parser(ParamParser) +.set_attr("FGetAttrDict", ParamGetAttrDict) +.set_attr("FInferShape", Pool2DInferShape) .set_attr("FInferType", ElemwiseType<1, 1>) -.set_attr("FCorrectLayout", Pool2DCorrectLayout) +.set_attr("FCorrectLayout", Pool2DCorrectLayout) .set_attr("FTVMCompute", [](const NodeAttrs& attrs, const Array& inputs, const Array& out_info) { - const Pool2DParam& param = nnvm::get(attrs.parsed); + const AvgPool2DParam& param = nnvm::get(attrs.parsed); auto pool_size = ShapeToArray(param.pool_size); auto strides = ShapeToArray(param.strides); auto padding = ShapeToArray(param.padding); auto ceil_mode = param.ceil_mode; + auto count_include_pad = param.count_include_pad; Layout layout(param.layout); CHECK(layout.convertible(Layout("NCHW"))) @@ -214,7 +218,7 @@ NNVM_REGISTER_OP(avg_pool2d) return Array{ topi::nn::pool(inputs[0], pool_size, strides, padding, - topi::nn::kAvgPool, ceil_mode, layout.name())}; + topi::nn::kAvgPool, ceil_mode, layout.name(), count_include_pad)}; }) .set_num_outputs(1) .set_num_inputs(1) diff --git a/nnvm/tests/python/compiler/test_top_level2.py b/nnvm/tests/python/compiler/test_top_level2.py index 44767da4541f..f752b7d44271 100644 --- a/nnvm/tests/python/compiler/test_top_level2.py +++ b/nnvm/tests/python/compiler/test_top_level2.py @@ -141,6 +141,41 @@ def test_avg_pool2d(): np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5) +def test_avg_pool2d_no_count_pad(): + kh, kw = (4, 4) + sh, sw = (2, 2) + ph, pw = (2, 2) + + x = sym.Variable("x") + y = sym.avg_pool2d(x, pool_size=(kh, kw), strides=(sw, sw), padding=(ph, pw), + name="y", count_include_pad=False) + dtype = "float32" + n = 1 + (ic, ih, iw) = (3, 28, 28) + (oc, oh, ow) = (3, 15, 15) + + a_np = np.random.uniform(low=0.001, size=(n, ic, ih, iw)).astype(dtype) + pad_np = np.zeros(shape=(n, ic, ih+2*ph, iw+2*pw)).astype(dtype) + no_zero = (range(n), range(ic), (range(ph, ih+ph)), (range(pw, iw+pw))) + pad_np[np.ix_(*no_zero)] = a_np + b_np = np.zeros(shape=(n, oc, oh, ow)).astype(dtype) + + for i in range(oh): + for j in range(ow): + pad_count = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw] > 0, axis=(2,3)) + b_np[:,:,i,j] = np.sum(pad_np[:, :, i*sh:i*sh+kh, j*sw:j*sw+kw], + axis=(2,3)) / np.maximum(pad_count, 1) + b_np = np.maximum(b_np, 0.0) + shape_dict = {"x": (n, ic, ih, iw)} + for target, ctx in ctx_list(): + graph, lib, _ = nnvm.compiler.build(y, target, shape_dict) + m = graph_runtime.create(graph, lib, ctx) + data = tvm.nd.array(a_np) + m.run(x=data) + out = m.get_output(0, tvm.nd.empty((n, oc, oh, ow), dtype)) + np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5) + + def test_global_max_pool2d(): x = sym.Variable("x") y = sym.global_max_pool2d(x, name="y") @@ -201,6 +236,7 @@ def test_upsampling(): test_conv2d_transpose() test_max_pool2d() test_avg_pool2d() + test_avg_pool2d_no_count_pad() test_global_max_pool2d() test_global_avg_pool2d() test_upsampling()