feat(imperative/python): support pooling mode "average" for avg pool2d module

GitOrigin-RevId: 9fe442129f

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__
+        )

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)