VGGNet

In [ ]:

from torch import nn

from torch import nn

In [ ]:

class VGGBlock(nn.Module):
    def __init__(self, num_convs, in_channels, out_channels):
        super(VGGBlock, self).__init__()
        layers = []
        for _ in range(num_convs):
            layers.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1))
            layers.append(nn.ReLU())
            in_channels = out_channels
        
        layers.append(nn.MaxPool2d(kernel_size=2, stride=2, padding=0))
        self.VGGBlock_layer = nn.Sequential(*layers)

    def forward(self, x):
        out = self.VGGBlock_layer(x)
        return out

class VGGBlock(nn.Module): def __init__(self, num_convs, in_channels, out_channels): super(VGGBlock, self).__init__() layers = [] for _ in range(num_convs): layers.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)) layers.append(nn.ReLU()) in_channels = out_channels layers.append(nn.MaxPool2d(kernel_size=2, stride=2, padding=0)) self.VGGBlock_layer = nn.Sequential(*layers) def forward(self, x): out = self.VGGBlock_layer(x) return out

In [ ]:

class VGGNet(nn.Module):
    def __init__(self, dropout=0.5):
        super(VGGNet, self).__init__()

        self.VGGBlock1 = VGGBlock(num_convs=2, in_channels=1, out_channels=64)
        self.VGGBlock2 = VGGBlock(num_convs=2, in_channels=64, out_channels=128)
        self.VGGBlock3 = VGGBlock(num_convs=3, in_channels=128, out_channels=256)
        self.VGGBlock4 = VGGBlock(num_convs=3, in_channels=256, out_channels=512)
        self.VGGBlock5 = VGGBlock(num_convs=3, in_channels=512, out_channels=512)

        self.FC = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(dropout),
            nn.Linear(512 * 2 * 2, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        out = self.VGGBlock1(x)
        out = self.VGGBlock2(out)
        out = self.VGGBlock3(out)
        out = self.VGGBlock4(out)
        out = self.VGGBlock5(out)
        out = self.FC(out)
        return out

class VGGNet(nn.Module): def __init__(self, dropout=0.5): super(VGGNet, self).__init__() self.VGGBlock1 = VGGBlock(num_convs=2, in_channels=1, out_channels=64) self.VGGBlock2 = VGGBlock(num_convs=2, in_channels=64, out_channels=128) self.VGGBlock3 = VGGBlock(num_convs=3, in_channels=128, out_channels=256) self.VGGBlock4 = VGGBlock(num_convs=3, in_channels=256, out_channels=512) self.VGGBlock5 = VGGBlock(num_convs=3, in_channels=512, out_channels=512) self.FC = nn.Sequential( nn.Flatten(), nn.Dropout(dropout), nn.Linear(512 * 2 * 2, 256), nn.ReLU(), nn.Linear(256, 256), nn.ReLU(), nn.Dropout(dropout), nn.Linear(256, 10) ) def forward(self, x): out = self.VGGBlock1(x) out = self.VGGBlock2(out) out = self.VGGBlock3(out) out = self.VGGBlock4(out) out = self.VGGBlock5(out) out = self.FC(out) return out