(구현) VERY DEEP CONVOLUTIONAL NETWORKS FOR LARGE-SCALE IMAGE RECOGNITION

VGG16 모델을 PyTorch로 구현해보았다. 각 파라미터 역할과 네트워크 구조를 이해할 수 있도록 주석을 달아두었다.

참고사이트

https://modulabs.co.kr/blog/vggnet_structure/

작은 커널, 더 큰 성과: VGGNet 아키텍처 탐구

https://github.com/msyim/VGG16/blob/master/VGG16.py

VGG16/VGG16.py at master · msyim/VGG16

import torch
import torch.nn as tnn
import torchvision.datasets as dsets
# 이미지 전처리 작업을 위함
import torchvision.transforms as transforms

# Compose는 여러 전처리 과정을 순차적으로 적용하기 위해 사용
transform = transforms.Compose([
    # 이미지에서 랜덤하게 자른 후, 해당 영역을 224x224 크기로 다시 맞춘다. ( VGG16 모델은 224x224 크기의 이미지를 입력으로 받음)
    transforms.RandomResizedCrop(224),
    # 이미지를 50% 확률로 좌우 반전 시킨다. (데이터 증강 기법)
    transforms.RandomHorizontalFlip(),
    # 이미지를 PyTorch 텐서로 변환한다.
    transforms.ToTensor()
    # 각 채널의 픽셀 값에 대해 평균과 표준편차를 기준으로 정규화한다.
    # mean은 ImageNet 데이터셋의 RGB 채널별 평균값
    # std는 채널별 표준편차
    transforms.Normalize(mean = [ 0.485, 0.456, 0.406],
                         std = [0.229, 0.224, 0.225]),
])

# 컨볼루션 레이어
def conv_layer(chann_in, chann_out, k_size, p_size):
    layer = tnn.Sequential(
        tnn.Conv2d(chann_in, chann_out, kernel_size=k_size, padding=p_size),
        tnn.ReLU()
    )
    return layer

# 컨볼루션 블록
# 여러 개의 컨볼루션 레이어를 쌓은 후 마지막에 맥스풀링을 적용하여 이미지 크기 절반으로 줄임
# in_list : (입력 채널, 출력 채널)
# out_list : (입력 채널, 출력 채널)
# k_list : 커널 크기
# p_list : 패딩 크기
# pooling_k, pooling_s : 풀링 커널 크기, 스트라이드
# ex) [3,64],[64,64], [3,3], [1,1], 2, 2
# (3, 64, 3, 1) -> 3 입력 채널에서 64 출력 채널로, 커널 크기 3, 패딩 1
def vgg_conv_block(in_list, out_list, k_list, p_list, pooling_k, pooling_s):
    layers = [ conv_layer(in_list[i], out_list[i], k_list[i], p_list[i]) for i in range(len(in_list))]
    layers += [ tnn.MaxPool2d(kernel_size=pooling_k, stride = pooling_s)]
    return tnn.Sequential(*layer)

# 완전 연결 레이어
def vgg_fc_layer(size_in, size_out):
    layer = tnn.Sequential(
        tnn.Linear(size_in, size_out),
        tnn.ReLU()
    )
    return layer

class VGG16(tnn.Module):
    def __init__(self, n_classes=1000):
        super(VGG16, self).__init__()
        # 첫 번째 컨볼루션 레이어는 3채널 이미지를 받아 64개의 출력 채널 생성
        # 두 번째 컨볼루션 레이어는 64개의 입력 채널을 받아 다시 64개의 출력 채널 생성
        # 마지막으로 MaxPooling 레이어가 2x2 풀링 커널을 사용하여 특성 맵의 크기 절반으로 줄임
        self.layer1 = vgg_conv_block([3,64],[64,64], [3,3], [1,1], 2, 2)
        self.layer2 = vgg_conv_block([64,128],[128,128],[3,3], [1,1],2, 2)
        self.layer3 = vgg_conv_block([128,256,256], [256,256,256], [3,3,3], [1,1,1],2,2)
        self.layer4 = vgg_conv_block([256,512,512], [512,512,512], [3,3,3], [1,1,1],2,2)
        self.layer5 = vgg_conv_block([512,512,512], [512,512,512], [3,3,3], [1,1,1],2,2)

        # FC layers
        self.layer6 = vgg_fc_layer(7*7*512,4096)
        self.layer7 = vgg_fc_layer(4096,4096)

        # Final layer
        self.layer8 = tnn.Linear(4096, n_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        # VGG16 모델에서 중간 단계의 특징 맵을 저장하는 변수
        out = self.layer5(out)
        # 특징 맵 평탄화
        out = out.view(out.size(0), -1)
        out = self.layer6(out)
        out = self.layer7(out)
        out = self.layer8(out)
        
        return out