SPP layer implementation in Pytorch

1

Guys I implemented Spatial Pyramid Pooling as follows but it seems to be so slow. Do you have any suggestion ?

class SPPLayer(nn.Module):

def __init__(self, num_levels, pool_type='max_pool'):
    super(SPPLayer, self).__init__()

    self.num_levels = num_levels
    self.pool_type = pool_type

def forward(self, x):
    bs, c, h, w = x.size()
    pooling_layers = []
    for i in range(self.num_levels):
        kernel_size = h // (2 ** i)
        if self.pool_type == 'max_pool':
            tensor = F.max_pool2d(x, kernel_size=kernel_size,
                                  stride=kernel_size).view(bs, -1)
        else:
            tensor = F.avg_pool2d(x, kernel_size=kernel_size,
                                  stride=kernel_size).view(bs, -1)
        pooling_layers.append(tensor)
    x = th.cat(pooling_layers, dim=-1)
    return x
2 Likes
4

According to the original paper of SPP-net, I think the kernel_size and stride for each pooling should be set as follows:

level = 2 ** i
kernel_size = (math.ceil(h / level), math.ceil(w / level))
stride = (math.floor(h/ level), math.floor(w / level))

1 Like