Applies a 2D soft pooling over an input signal composed of several input planes. See https://arxiv.org/abs/2101.00440

softpool.py

from torch import nn
from torch.nn import functional as F


class SoftPool2d(nn.Module):
    """Applies a 2D soft pooling over an input signal composed of several
    input planes. See https://arxiv.org/abs/2101.00440"""

    def __init__(self, kernel_size, ceil_mode=False, temperature=1.):
        super().__init__()
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)
        self.kernel_size = kernel_size
        self.ceil_mode = ceil_mode
        self.temperature = temperature

    def extra_repr(self):
        return f'kernel_size={self.kernel_size}, ' \
               f'ceil_mode={self.ceil_mode}, ' \
               f'temperature={self.temperature:g}'

    def forward(self, input):
        kh, kw = self.kernel_size
        h, w = input.shape[2:]
        pad_h = h % kh
        pad_w = w % kw
        if self.ceil_mode:
            input = F.pad(input, (0, pad_w, 0, pad_h))
        else:
            input = input[..., :h-pad_h, :w-pad_w]
        n, c, h, w = input.shape
        input = input.view([n, c, h//kh, kh, w//kw, kw]).movedim(3, 4).flatten(4)
        input = input * F.softmax(input / self.temperature, dim=4)
        return input.sum(dim=4)