| 9 | |
| 10 | def test_avg_pool2d(): |
| 11 | def test_func( |
| 12 | batch_size, |
| 13 | in_channels, |
| 14 | out_channels, |
| 15 | in_height, |
| 16 | in_width, |
| 17 | kernel_size, |
| 18 | stride, |
| 19 | padding, |
| 20 | ): |
| 21 | pool = AvgPool2d(kernel_size, stride=stride, padding=padding, mode="average") |
| 22 | inp = np.random.normal( |
| 23 | size=(batch_size, in_channels, in_height, in_width) |
| 24 | ).astype(np.float32) |
| 25 | out_height = (in_height + padding * 2 - kernel_size) // stride + 1 |
| 26 | out_width = (in_width + padding * 2 - kernel_size) // stride + 1 |
| 27 | out = pool(tensor(inp)) |
| 28 | inp = np.pad(inp, ((0, 0), (0, 0), (padding, padding), (padding, padding))) |
| 29 | expected = np.zeros( |
| 30 | (batch_size, out_channels, out_height, out_width), dtype=np.float32, |
| 31 | ) |
| 32 | for n, c, oh, ow in itertools.product( |
| 33 | *map(range, [batch_size, out_channels, out_height, out_width]) |
| 34 | ): |
| 35 | ih, iw = oh * stride, ow * stride |
| 36 | expected[n, c, oh, ow] = np.sum( |
| 37 | inp[n, c, ih : ih + kernel_size, iw : iw + kernel_size,] |
| 38 | ) / (kernel_size * kernel_size) |
| 39 | np.testing.assert_almost_equal(out.numpy(), expected, 1e-5) |
| 40 | |
| 41 | test_func(10, 4, 4, 5, 5, 2, 2, 1) |
| 42 | test_func(10, 4, 4, 6, 6, 2, 2, 0) |