Spp-Net Pytorch

基于Pytorch复现Spp-Net

Pytorch：0.4  ； Python：3.5。较tensorflow易实现。

'''

sppnet base on vgg16
input size random ; but batch size need set to be 1 , and we don't use 'batch_size' .
'''

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

class SppNet(nn.Module):
    def __init__(self, batch_size=1, out_pool_size=[1,2,4], class_number=2):
        super().__init__()
        # use already written network , eg vgg16
        vgg = torchvision.models.vgg16(pretrained=False).features[:-1]
        self.out_pool_size = out_pool_size
        self.batch_size=batch_size
        # encoder
        self.encoder = vgg
        # spp if spp is a class , so create network ,it appear (spp)
        self.spp = self.make_spp(batch_size=batch_size, out_pool_size=out_pool_size)
        # FC
        sum0=0
        for i in out_pool_size:
            sum0+=i**2
        self.fc = nn.Sequential(nn.Linear(512*sum0,1024), nn.ReLU(inplace=True))
        self.score = nn.Linear(1024,class_number)
        
    def make_spp(self, batch_size=1, out_pool_size=[1,2,4]):
        func=[]
        for i in range(len(out_pool_size)):
            func.append(nn.AdaptiveAvgPool2d(output_size=(out_pool_size[i],out_pool_size[i])))
        return func
    
    def forward(self, x):
        assert x.shape[0] == 1 , 'batch size need to set to be 1'
        encoder = self.encoder(x)
        spp=[]
        for i in range(len(self.out_pool_size)):
            spp.append( self.spp[i](encoder).view(self.batch_size,-1) )
        fc = self.fc(torch.cat(spp,dim=1))
        score = self.score(fc)
        return score
    
''' or another '''  
         
class SppNet1(nn.Module):
    def __init__(self, batch_size=1, out_pool_size=[1,2,4], class_number=2):
        super().__init__()
        # use already written network , eg vgg16
        vgg = torchvision.models.vgg16(pretrained=False).features[:-1]
        self.out_pool_size = out_pool_size
        self.batch_size=batch_size
        # encoder
        self.encoder = vgg
        # FC
        sum0=0
        for i in out_pool_size:
            sum0+=i**2
        self.fc = nn.Sequential(nn.Linear(512*sum0,1024), nn.ReLU(inplace=True))
        self.score = nn.Linear(1024,class_number)
    
    def forward(self, x):
        assert x.shape[0] == 1 , 'batch size need to set to be 1'
        encoder = self.encoder(x)
        spp=[]
        for i in range(len(self.out_pool_size)):
            spp.append( F.adaptive_avg_pool2d(encoder,output_size=(self.out_pool_size[i],self.out_pool_size[i])).view(self.batch_size,-1) )
        fc = self.fc(torch.cat(spp,dim=1))
        score = self.score(fc)
        return score     
        
spp=SppNet(class_number=2)

print(spp)