卷积神经网络模型之ZFNet模型实现

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

class ZFNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1=nn.Conv2d(in_channels=3,out_channels=96,kernel_size=7,stride=2,padding=1) #[110,110,96]
        self.maxpool1=nn.MaxPool2d(kernel_size=3,stride=2,padding=1) #[55,55,96]
        
        self.conv2=nn.Conv2d(in_channels=96,out_channels=256,kernel_size=5,stride=2)#[25,25,256]
        self.maxpool2=nn.MaxPool2d(kernel_size=3,stride=2,padding=1)#[13,13,256]
        
        self.conv3=nn.Conv2d(in_channels=256,out_channels=384,kernel_size=3,stride=1,padding=1)
        
        self.conv4=nn.Conv2d(in_channels=384,out_channels=384,kernel_size=3,stride=1,padding=1)
       
        self.conv5=nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,stride=1,padding=1)
        self.maxpool3=nn.MaxPool2d(kernel_size=3,stride=2)
        
        self.fc1=nn.Linear(in_features=6*6*256,out_features=4096)
        
        self.fc2=nn.Linear(in_features=4096,out_features=4096)
        
        self.fc3=nn.Linear(in_features=4096,out_features=1000)
        
        
    def forward(self,x):
        x=F.relu(self.conv1(x))
        print("conv1:",x.shape)
        x=self.maxpool1(x)
        print("maxpool1:",x.shape)
        
        x=F.relu(self.conv2(x))
        print("conv2:",x.shape)
        x=self.maxpool2(x)
        print("maxpool2:",x.shape)
        
        x=F.relu(self.conv3(x))
        print("conv3:",x.shape)
        x=F.relu(self.conv4(x))
        print("conv4:",x.shape)
        x=self.maxpool3(F.relu(self.conv5(x)))
        print("maxpool3:",x.shape)
        
        x=x.view(-1,6*6*256)
        print(x.shape)
        x=F.relu(self.fc1(x))
        print("fc1",x.shape)
        x=F.relu(self.fc2(x))
        print("fc2",x.shape)
        x=self.fc3(x)
        print("fc3",x.shape)
        return x

model=ZFNet()
model.cuda()
print(model)

from torch.autograd import Variable
x=torch.randn(1,3,224,224)
x=Variable(x).cuda()
pred=model(x)

print(pred.shape)