diff --git a/imperative/python/megengine/module/pooling.py b/imperative/python/megengine/module/pooling.py index c8ff81be..5bbecf60 100644 --- a/imperative/python/megengine/module/pooling.py +++ b/imperative/python/megengine/module/pooling.py @@ -103,6 +103,9 @@ class AvgPool2d(_PoolNd): :param kernel_size: the size of the window. :param stride: the stride of the window. Default value is kernel_size。 :param padding: implicit zero padding to be added on both sides. + :param mode: whether to count padding values. "average" mode will do counting and + "average_count_exclude_padding" mode won't do counting. + Default: "average_count_exclude_padding" Examples: @@ -126,5 +129,21 @@ class AvgPool2d(_PoolNd): """ + def __init__( + self, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]] = None, + padding: Union[int, Tuple[int, int]] = 0, + mode: str = "average_count_exclude_padding", + **kwargs + ): + super(AvgPool2d, self).__init__(kernel_size, stride, padding, **kwargs) + self.mode = mode + def forward(self, inp): - return avg_pool2d(inp, self.kernel_size, self.stride, self.padding) + return avg_pool2d(inp, self.kernel_size, self.stride, self.padding, self.mode) + + def _module_info_string(self) -> str: + return "kernel_size={kernel_size}, stride={stride}, padding={padding}, mode={mode}".format( + **self.__dict__ + ) diff --git a/imperative/python/test/unit/module/test_pool2d.py b/imperative/python/test/unit/module/test_pool2d.py new file mode 100644 index 00000000..9f06469f --- /dev/null +++ b/imperative/python/test/unit/module/test_pool2d.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +import itertools + +import numpy as np + +from megengine import Parameter, tensor +from megengine.module import AvgPool2d, MaxPool2d + + +def test_avg_pool2d(): + def test_func( + batch_size, + in_channels, + out_channels, + in_height, + in_width, + kernel_size, + stride, + padding, + ): + pool = AvgPool2d(kernel_size, stride=stride, padding=padding, mode="average") + inp = np.random.normal( + size=(batch_size, in_channels, in_height, in_width) + ).astype(np.float32) + out_height = (in_height + padding * 2 - kernel_size) // stride + 1 + out_width = (in_width + padding * 2 - kernel_size) // stride + 1 + out = pool(tensor(inp)) + inp = np.pad(inp, ((0, 0), (0, 0), (padding, padding), (padding, padding))) + expected = np.zeros( + (batch_size, out_channels, out_height, out_width), dtype=np.float32, + ) + for n, c, oh, ow in itertools.product( + *map(range, [batch_size, out_channels, out_height, out_width]) + ): + ih, iw = oh * stride, ow * stride + expected[n, c, oh, ow] = np.sum( + inp[n, c, ih : ih + kernel_size, iw : iw + kernel_size,] + ) / (kernel_size * kernel_size) + np.testing.assert_almost_equal(out.numpy(), expected, 1e-5) + + test_func(10, 4, 4, 5, 5, 2, 2, 1) + test_func(10, 4, 4, 6, 6, 2, 2, 0) + test_func(10, 16, 16, 14, 14, 2, 2, 0)