|
- # -*- 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)
|
