[논문 구현] Pytorch 기반 VGG 구현 및 CIFAR10 실험

---

 

Pytorch로 VGG 논문 구현하기

 

논문 링크 : https://arxiv.org/abs/1409.1556

Very Deep Convolutional Networks for Large-Scale Image Recognition

 

해당 논문은 리뷰하기엔 내용이 조금 짧아서 바로 구현을 시작하기로 했다.

어찌보면, 논문 구현의 Hello World! 와도 같은 느낌의 VGG이다.

 

 

라이브러리 import

---

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

 

논문 제시 Architecture

---

 

VGG의 Architecture는 위와 같다. 이중에서도 특히 D와 E가 흔히 알려진 VGG-16, VGG-19 모델이다.

VGG는 아직까지도 간단한 구조로 Feature Map등을 추출하는 용도의 backbone으로 사용되고 있다.

 

각각의 Conv 이후에는 ReLU가 존재하지만, 표에서는 빠져있다고 논문에 명시되어 있으므로 이를 참고하여 구현해보자.

 

Make Layer 

---

VGG는 매우 간단한 구조로, 특정 Block 형태가 반복되기보다는 그냥 똑같히 생긴 Conv 레이어와 Maxpool이 반복되는 구조이다. 따라서, ResNet과 같이 Block 단위로 정의할 필요 없이, 다소 무식한 방법으로 레이어를 만들어 나가는 것이 용이하다.

 

cfg = {
    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}

# A : VGG-11
# B : VGG-13
# D : VGG-16
# E : VGG-19

def make_layer(config):
    layers = []
    in_planes = 3
    for value in config:
        if value == "M":
            layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
        else:
            layers.append(nn.Conv2d(in_planes, value, kernel_size=3, padding=1))
            layers.append(nn.ReLU())
            in_planes = value
    return nn.Sequential(*layers)

 

 

VGG 모델 정의

---

class VGG(nn.Module):
    def __init__(self, config, num_classes=1000, cifar=False):
        super(VGG, self).__init__()
        self.features = make_layer(config)
        
        # ImageNet
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, num_classes)  
        )
        # CIFAR-10
        if cifar:
            self.classifier = nn.Sequential(
                nn.Dropout(0.5),
                nn.Linear(512, 512),
                nn.ReLU(True),
                nn.Dropout(0.5),
                nn.Linear(512, 512),
                nn.ReLU(True),
                nn.Linear(512, 10)  
            ) 
        
    def forward(self, x):
        out = self.features(x)
        out = torch.flatten(out,1)
        out = self.classifier(out)
        return out

 

Classifier의 경우 항상 동일한 조건이므로, make layer 이후에 연산을 진행해주기만 하면 된다.

이 때, ImageNet 기준으로는 Linear 4096, 1000을 마지막 연산으로 수행하지만, CIFAR-10의 경우 이미지 크기가 다르고, 최종 n_classes가 10이므로, 512로 변경하도록 했다.

 

VGG Architecture 생성

---

def VGG11(cifar=False):
    return VGG(config = cfg['A'], cifar = cifar)

def VGG13(cifar=False):
    return VGG(config = cfg['B'], cifar = cifar)

def VGG16(cifar=False):
    return VGG(config = cfg['D'], cifar = cifar)

def VGG19(cifar=False):
    return VGG(config = cfg['E'], cifar = cifar)

 

논문 구현을 마치고 가장 뿌듯한 순간이다.

아키텍쳐 생성을 미리 정의해준 cfg을 통해 진행해준다.

 

이제 직접 구현한 VGG를 사용할 수 있다.

해당 구현 코드와 CIFAR-10 학습 코드는 Jupyter Notebook을 통해 Github에 올려두었으니, 

Colab 무료 GPU 서버를 통해 실습해볼 수 있다.

 

Github : https://github.com/CryptoSalamander/pytorch_paper_implementation/tree/master/vgg