본문 바로가기
Computer Vision

VGGNet PyTorch Code Implementation

by 아르카눔 2024. 4. 2.

VGGNet을 Pytorch를 활용하여 구현하고자 한다.

https://arsetstudium.tistory.com/29에서 공부한 내용을 토대로 구현보면 아래와 같다.

 

우선 Top-5 error rate 성능이 가장 낮았던, 성능이 가장 좋았던 16 layers와 19 layers를 구현하는 편이 좋다.

그중에서도 16 layers가 파라미터 수가 더 적어서 효율적이기 때문에 16 layers 구조를 채택해서 구현하고자 한다.

 

VGGNet은 max pooling을 경계로 해서 features extraction의 convolutional layers를 5개의 block으로 나눠 볼수도 있다.

하지만 여기서는 구체적으로 block의 class나 sequential로 나누지는 않았고 주석으로만 표기했다.

 

 

class VGGNet(nn.Module):
    def __init__(self, num_classes = 1000, dropout = 0.5):
        super().__init__()
        self.features = nn.Sequential(
            # Extracting Fetures Part
            # First Block
            nn.Conv2d(3, 64, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size = 3, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2, stride = 2),
            # Second Block
            nn.Conv2d(64, 128, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(128, 128, kernel_size = 3, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2, stride = 2),
            # Third Block
            nn.Conv2d(128, 256, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size = 3, padding = 1),
            nn.ReLU(),
            nn.Conv2d(256, 256, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2, stride = 2),
            # Fourth Block
            nn.Conv2d(256, 512, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size = 3, padding = 1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2, stride = 2),
            # Fifth Block
            nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size = 3, padding = 1),
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size = 3, stride = 1, padding = 1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2, stride = 2),                 

        )

        self.classifier = nn.Sequential(
            # Classification Part
            nn.Dropout(p=dropout),
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Linear(4096, num_classes),
            nn.Softmax(dim=0)
        )

    def forward(self, x):
        y = self.features(x)
        y = torch.flatten(y, start_dim = 1)
        y = self.classifier(y)
        return y

 

AlexNet과 마찬가지로 classifier로 넘어가기 전에 flatten을 해야하는데 in_features를 알아야 한다.

self.features를 모두 계산하여 224 x 224 image가 어떻게 변하나 확인할 수도 있지만,

모델을 선언한 다음 input data에 self.features만 적용하여 확인할 수도 있다.

 

a = torch.rand(5, 3, 224, 224)
b = model.features(a)
b.shape

>> torch.Size([5, 512, 7, 7])

 

5개의 batch size를 가지는 데이터의 224 x 224 RGB 데이터를 넣었다고 가정해보았다.

model.features를 통해 도출된 결과물의 dimension은 5 x 512 x 7 x 7이다.

이때 batch size를 제외한 나머지 항목을 모두 곱하면 classifier의 첫번째 FC layer에 들어갈 in_features의 크기를 구할 수 있다. 여기서는 그 값이 25088이다.

 

VGGNet-16 from torchvision

 

torchvision에서 제공하는 VGGNet-16 코드를 통해 직접 구현한 코드를 비교하고자 한다.

 

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

 

AlexNet과 마찬가지로 중간에 AdaptiveAvgPool2d가 있다.

또한 다른점이 있다면 AlexNet은 flatten한 다음에 dropout이 있었는데,

VGGNet에서는 1번째 FC layer와 2번재 FC layer의 사이에서 처음 dropout이 등장한다.

이 차이는 어디서 오는지는 추가적으로 알아봐야겠다.

 

 

AlexNet에서는 구체적인 학습까지 진행했지만 여기서는 자세한 프로세스 자체를 똑같은 image classification이라서 생략한다.