Tensorflow like 'SAME' wrapper for 2D average pooling
| 19 | |
| 20 | |
| 21 | class AvgPool2dSame(nn.AvgPool2d): |
| 22 | """ Tensorflow like 'SAME' wrapper for 2D average pooling |
| 23 | """ |
| 24 | def __init__(self, kernel_size: int, stride=None, padding=0, ceil_mode=False, count_include_pad=True): |
| 25 | kernel_size = to_2tuple(kernel_size) |
| 26 | stride = to_2tuple(stride) |
| 27 | super(AvgPool2dSame, self).__init__(kernel_size, stride, (0, 0), ceil_mode, count_include_pad) |
| 28 | |
| 29 | def forward(self, x): |
| 30 | x = pad_same(x, self.kernel_size, self.stride) |
| 31 | return F.avg_pool2d( |
| 32 | x, self.kernel_size, self.stride, self.padding, self.ceil_mode, self.count_include_pad) |
| 33 | |
| 34 | |
| 35 | def max_pool2d_same( |