pytorch实现vgg11模型

VGG11

这里实现最左侧的A，也就是原始的vgg11。

import torch
import torch.nn as nn


class VGG11(nn.Module):
    def __init__(self, input_channel):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(in_channels=input_channel, out_channels=64,  # [B,64,224,224]
                      kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2),  # [B,64,112,112]
            nn.Conv2d(in_channels=64, out_channels=128,  # [B,128,112,112]
                      kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2),  # [B,128,56,56]
            # [B,256,56,56]
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            # [B,256,56,56]
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2),  # [B,256,28,28]
            # [B,512,56,56]
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            # [B,512,56,56]
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2),  # [B,512,14,14]
            # [B,512,28,28]
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            # [B,512,56,56]
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=(3, 3), padding=(3 // 2, 3 // 2)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=2)  # [B,512,7,7]
        )

        self.fc_layers=nn.Sequential(
            nn.Linear(512*7*7,4096),
            nn.ReLU(),
            nn.Linear(4096,4096),
            nn.ReLU(),
            nn.Linear(4096,1000),
            nn.Softmax(dim=-1)
        )

    def forward(self, x):
        conv_out = self.conv_layers(x)
        fc_out=self.fc_layers(conv_out.view(conv_out.size(0),-1))
        return fc_out

if __name__ == '__main__':
    x = torch.randn(size=(8, 3, 224, 224))  # [N,C,W,H]
    vgg11 = VGG11(input_channel=3)
    conv_out = vgg11(x)
    print(conv_out.size())