利用pytorch复现spatial pyramid pooling层

sppnet不讲了，懒得写。。。直接上代码

 from math import floor, ceil
 import torch
 import torch.nn as nn
 import torch.nn.functional as F

 class SpatialPyramidPooling2d(nn.Module):
     r"""apply spatial pyramid pooling over a 4d input(a mini-batch of 2d inputs
     with additional channel dimension) as described in the paper
     'Spatial Pyramid Pooling in deep convolutional Networks for visual recognition'
     Args:
         num_level:
         pool_type: max_pool, avg_pool, Default:max_pool
     By the way, the target output size is num_grid:
         num_grid = 0
         for i in range num_level:
             num_grid += (i + 1) * (i + 1)
         num_grid = num_grid * channels # channels is the channel dimension of input data
     examples:
         >>> input = torch.randn((1,3,32,32), dtype=torch.float32)
         >>> net = torch.nn.Sequential(nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3,stride=1),\
                                       nn.ReLU(),\
                                       SpatialPyramidPooling2d(num_level=2,pool_type='avg_pool'),\
                                       nn.Linear(32 * (1*1 + 2*2), 10))
         >>> output = net(input)
     """

     def __init__(self, num_level, pool_type='max_pool'):
         super(SpatialPyramidPooling2d, self).__init__()
         self.num_level = num_level
         self.pool_type = pool_type

     def forward(self, x):
         N, C, H, W = x.size()
         for i in range(self.num_level):
             level = i + 1
             kernel_size = (ceil(H / level), ceil(W / level))
             stride = (ceil(H / level), ceil(W / level))
             padding = (floor((kernel_size[0] * level - H + 1) / 2), floor((kernel_size[1] * level - W + 1) / 2))

             if self.pool_type == 'max_pool':
                 tensor = (F.max_pool2d(x, kernel_size=kernel_size, stride=stride, padding=padding)).view(N, -1)
             else:
                 tensor = (F.avg_pool2d(x, kernel_size=kernel_size, stride=stride, padding=padding)).view(N, -1)

             if i == 0:
                 res = tensor
             else:
                 res = torch.cat((res, tensor), 1)
         return res
     def __repr__(self):
         return self.__class__.__name__ + '(' \
             + 'num_level = ' + str(self.num_level) \
             + ', pool_type = ' + str(self.pool_type) + ')'

 class SPPNet(nn.Module):
     def __init__(self, num_level=3, pool_type='max_pool'):
         super(SPPNet,self).__init__()
         self.num_level = num_level
         self.pool_type = pool_type
         self.feature = nn.Sequential(nn.Conv2d(3,64,3),\
                                     nn.ReLU(),\
                                     nn.MaxPool2d(2),\
                                     nn.Conv2d(64,64,3),\
                                     nn.ReLU())
         self.num_grid = self._cal_num_grids(num_level)
         self.spp_layer = SpatialPyramidPooling2d(num_level)
         self.linear = nn.Sequential(nn.Linear(self.num_grid * 64, 512),\
                                     nn.Linear(512, 10))
     def _cal_num_grids(self, level):
         count = 0
         for i in range(level):
             count += (i + 1) * (i + 1)
         return count

     def forward(self, x):
         x = self.feature(x)
         x = self.spp_layer(x)
         print(x.size())
         x = self.linear(x)
         return x

 if __name__ == '__main__':
     a = torch.rand((1,3,64,64))
     net = SPPNet()
     output = net(a)
     print(output)