PyTorch——模型搭建——VGG（一）

模型搭建

0、VGG模型

1、torchvision自带的VGG模型

import torch
import torchvision
from torchsummary import summary

model = torchvision.models.vgg13(num_classes=8).cuda()
summary(model, (3, 224, 224))

2、自己搭建

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

class VGG(nn.Module):

    def __init__(self, arch, num_classes=1000):
        super(VGG, self).__init__()
        self.in_channels = 3
        self.conv3_64 = self._make_layer(64, arch[0])
        self.conv3_128 = self._make_layer(128, arch[1])
        self.conv3_256 = self._make_layer(256, arch[2])
        self.conv3_512a = self._make_layer(512, arch[3])
        self.conv3_512b = self._make_layer(512, arch[4])
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(7*7*512, 4096)
        self.bn1 = nn.BatchNorm1d(4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)

    def _make_layer(self, channels, num):
        layers = []
        for i in range(num):
            layers.append(nn.Conv2d(self.in_channels, channels, 3, stride=1, padding=1, bias=False))
            layers.append(nn.BatchNorm2d(channels))
            layers.append(nn.ReLU())
            self.in_channels = channels
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv3_64(x)
        x = F.max_pool2d(x, 2)
        x = self.conv3_128(x)
        x = F.max_pool2d(x, 2)
        x = self.conv3_256(x)
        x = F.max_pool2d(x, 2)
        x = self.conv3_512a(x)
        x = F.max_pool2d(x, 2)
        x = self.conv3_512b(x)
        x = F.max_pool2d(x, 2)
        #x = x.view(x.size(0), -1)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.fc3(x)
        return x

def VGG_11():
    return VGG([1,1,2,2,2], num_classes=6)

def VGG_13():
    return VGG([2,2,2,2,2], num_classes=7)

def VGG_16():
    return VGG([2,2,3,3,3], num_classes=8)

def VGG_19():
    return VGG([2,2,4,4,4], num_classes=9)

net = VGG_16().cuda()
summary(net, (3, 224, 224))