[논문 구현] VGG16 구현하기 (By Colab, PyTorch)

웅장한 VGG16

 

이번 포스팅은 VGG 논문 구현이다. 만약 VGG 논문 혹은 논문 리뷰를 읽지 않았다면 읽고 보는 것을 추천한다.

 

https://stydy-sturdy.tistory.com/7

[논문 리뷰] Very Deep Convolutional Networks For Large-Scale Image Recognition (VGG) - 구조, Training 그리고 Test

 

포스팅의 시작을 VGG16 모델의 summary로 시작해보았다. VGG 모델을 만들고 난 후 LeNet-5를 다시보니 LeNet-5가 참 쪼꼬미 모델이라는 것이 다시 느껴진다. 

 

 

저번 시간 구현했던 LeNet-5의 summary

 

우선 VGG16의 구현의 느낀점부터 간략히 말하면 VGG16에는 parameter도 너무 많고 학습도 잘 되지 않는다. 초반에는 내가 모델 설계를 잘못한 건가 싶었지만 VGG 구현을 도전한 다른 블로그들을 봐도 잘 학습이 되지 않는 것이 사실인 것 같다. (Colab이 느린건지도..)

 

본격적으로 VGG 구현에 들어가기 앞서, VGG의 구조에 대해 복습해보자.

 

VGG의 구조

 

앞서 논문 리뷰에서도 말했지만 개인적으로 VGG에서는 위 표의 D와 E가 핵심이라고 생각한다. 각각 VGG16, VGG19인데, 뒤에 붙은 숫자는 layer의 개수를 의미한다. 다시 한번 모델의 구조를 살펴보자.

• 인풋 데이터 : 224*224 RGB 이미지
• 전처리 : training set의 mean RGB 값을 각 픽셀에서 빼주었다.
• Convolution stride : fixed to 1 pixel
• 5개의 max-pooling layer : 2*2 pixel에 대해 적용, stride = 2 (non-overlapping)
• Fully connected layer 1 : 4096 채널, ReLu
• Fully connected layer 2 : 4096 채널, ReLu
• Fully connected layer 3 : 1000 채널 (이 때 당시 이미지넷 데이터셋의 class의 수), Softmax

이것은 layer수에 상관없이 공통적인 부분이고, VGG16의 경우 3*3 필터를 13층, VGG19의 경우 3*3필터를 16층 사용하여 학습시킨다. 왜 3*3 필터를 사용하는지에 대해서는 저번 시간 충분히 설명했다 생각하고 넘어가도록 하겠다. (요약하면 3*3필터가 깊이가 깊어지고 비선형성이 증가해 이로운 점이 많다는 것이었다.)

 

이번 시간에는 특별히 VGG16의 구현에 대해서만 다루고, 다음 시간에는 더 똑똑한(?) 방법을 사용하여 VGG 모델들을 구현할 수 있는 코드를 소개하도록 하겠다.

 

VGG16의 구현을 위해서 정의해야하는 것들은 다음과 같다.

1. 모델의 정확도 측정 함수
2. 손실 plotting 함수
3. Train 함수
4. Validate 함수
5. Training_loop 함수
6. Class VGG16
7. Hyperparameters

LeNet-5의 구현과 아주 비슷한 방향으로 모델 설계를 해보았다. 개인적으로 생각되는 LeNet-5와 구별되는 점은 이미지의 전처리 과정이다. 논문에서  training set의 mean RGB 값을 각 픽셀에서 빼주었다고 했기 때문에 이 부분을 구현하는 것이 의미 있어 보인다. 추가적으로 LeNet-5과 마찬가지로 데이터셋을 불러오고 잘 저장이 되었나 확인하는 과정도 진행 해야한다.

 

논문에서는 이미지넷이라는 데이터셋을 사용하였지만, 이미지넷 데이터셋은 클래스가 무려 1000개이다. 나의 작은 노트북으로 이런 어마어마한 데이터셋을 모두 학습하는 것을 기대하긴 어렵기 때문에 다른 데이터셋을 선택하였다.

 

STL-10 Dataset 

 

https://cs.stanford.edu/~acoates/stl10/

STL-10 dataset

 

'The STL-10 dataset is an image recognition dataset for developing unsupervised feature learning, deep learning, self-taught learning algorithms. ' 

 

10개의 클래스가 있고, 컬러 이미지의 해상도는 96X96이다. train data의 수는 5000개이고, test data의 수는 8000개라고 한다.

 

이제 본격적으로 VGG16을 만들어보자! 1~5번까지는 LeNet-5와 거의 동일하다.

 

1. 모듈 가져오기

from google.colab import drive
drive.mount('VGG')

import numpy as np
from datetime import datetime
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
from pytz import timezone
import os

2. Hyperparameters

Random_seed = 42 #그냥 정하는 것
Learning_rate = 0.0001 # 학습률. 논문에서는 0.01로 시작하여 감소시키지만 우선 0.0001로 해보자
Batch_size = 64 # batch size. 이 단위로 학습이 된다.
N_epochs = 50 # epoch 횟수 전체 데이터가 50번 학습된다는 뜻.
Img_size = 224 # 224*224 RGB 이미지가 들어갈 것이다.
N_classes = 10 # 데이터셋의 10개의 output class가 있다.

 

Hyperparameter 설정 부분에서 시행 착오를 많이 겪었다. 시행 착오와 관련해서는 학습 부분에 한번 더 언급이 될 것이다.

 

논문에서 0.01로 시작해서 학습률을 점점 감소 시킨다고 나와있었지만, 그냥 0.01로 설정한 후 학습을 진행시켜봤더니 50에포크까지 정확도가 10%였다. 답이 없는 상황에서 학습률을 낮추었다. 0.001도 학습이 느려 0.0001로 설정한 후 진행하였다.

3. get_accuracy 함수 (모델의 정확도 측정)

# 모델의 정확도를 얻는 함수를 먼저 정의하자

def get_accuracy(model, data_loader, device) :
  correct_pred = 0
  n = 0
  with torch.no_grad() :
    model.eval() #batch nomalization, drop out과 같은거 없이! 모델이 평가모드로 전환
    for X, y_true in data_loader : #데이터 셋에 있는 인풋
      X = X.to(device)
      y_true = y_true.to(device)
      _, y_prob = model(X) #y_prob은 후에 나올 LeNet-5 모델에서 소프트맥스 함수를 통과한 클래스 확률
      _, predicted_labels = torch.max(y_prob,1) #torch.max(y_prob,1) = 열 중에서 가장 높은 값을 뽑아준다. 그럼 그게 예측 라벨이 되겠죠?

      n += y_true.size(0)
      correct_pred += (predicted_labels == y_true).sum()
  return correct_pred.float() / n # 정확히 맞춘것 / 총 개수

4. 손실 plotting 함수

def plot_loss(train_loss, val_loss) :
  plt.style.use('grayscale')
  train_loss = np.array(train_loss)
  val_loss = np.array(val_loss)
  fig , ax = plt.subplots(1,1,figsize = (8,4.5))
  ax.plot(train_loss, color = 'green' , label = 'Training Loss')
  ax.plot(val_loss, color = 'red' , label = 'Validation Loss')
  ax.set(title = 'Loss Over Epochs' , xlabel = 'EPOCH' , ylabel = 'LOSS')
  ax.legend()
  fig.show()
  plt.style.use('default')

5. Train 함수

def train(train_loader, model, criterion, optimizer, device) :
  model.train() #모델을 학습 모드로 설정
  running_loss = 0 # 초기값 0으로 설정
  for X, y_true in train_loader:
    optimizer.zero_grad() #역전파시 효과적으로 학습되기 위해 설정 매번 세팅되어야함
    X = X.to(device)
    y_true = y_true.to(device)
    y_hat, _ = model(X)
    loss = criterion(y_hat,y_true) #loss를 구함
    running_loss += loss.item() * X.size(0) #사이즈를 곱해줘서 전체적인 running loss를 구함
    loss.backward() #역전파
    optimizer.step() #Gradient descent
  epoch_loss = running_loss / len(train_loader.dataset)
  return model , optimizer, epoch_loss

6. Validate 함수

def validate(valid_loader, model, criterion, device):
    model.eval()
    running_loss = 0

    for X, y_true in valid_loader:
        X = X.to(device)
        y_true = y_true.to(device)

        # 순전파와 손실 기록하기
        y_hat, _ = model(X) #소프트 맥스 당하기 전 !
        loss = criterion(y_hat, y_true)
        running_loss += loss.item() * X.size(0)

    epoch_loss = running_loss / len(valid_loader.dataset)
    return model, epoch_loss

7. Training_loop 함수

def training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs, device, print_every=1):

    # metrics를 저장하기 위한 객체 설정
    best_loss = 1e10
    train_losses = []
    valid_losses = []

    # model 학습하기
    for epoch in range(0, epochs):

        # training
        model, optimizer, train_loss = train(train_loader, model, criterion, optimizer, device)
        train_losses.append(train_loss)

        # validation
        with torch.no_grad():
            model, valid_loss = validate(valid_loader, model, criterion, device)
            valid_losses.append(valid_loss)

        if epoch % print_every == (print_every - 1):

            train_acc = get_accuracy(model, train_loader, device=device)
            valid_acc = get_accuracy(model, valid_loader, device=device)

            print(datetime.now(timezone('Asia/Seoul')).time().replace(microsecond=0),'--- ',
                  f'Epoch: {epoch}\t'
                  f'Train loss: {train_loss:.4f}\t'
                  f'Valid loss: {valid_loss:.4f}\t'
                  f'Train accuracy: {100 * train_acc:.2f}\t'
                  f'Valid accuracy: {100 * valid_acc:.2f}')

    plot_loss(train_losses, valid_losses)

    return model, optimizer, (train_losses, valid_losses)

8. Class VGG16

VGG16

 

위 그림을 보며 코드를 따라오면 더 이해가 잘 될 것이다. 개인적으로는 Dense 위의 7*7*512가 중요하다고 생각한다.

왜 7*7*512 인풋이 classifier에 들어가는 것일까?

 

내가 이해한 바를 간략히 적어봤다. 결론적으로 MaxPooling을 5번 사용하였으므로 224*224 사이즈는 2**5 만큼 줄어들어 7*7이 되고, 마지막 채널의 수는 512이기 때문에 7*7*512 라는 값이 나오는 것이다.

 

class VGG16(nn.Module):

    def __init__(self, n_classes):
        super(VGG16, self).__init__()

        self.feature_extractor = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1,padding=1), #Bias default
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2), # default stride = kernel_size
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),            
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),  
            nn.MaxPool2d(kernel_size=2),   
            nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),  
            nn.MaxPool2d(kernel_size=2),  
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),   
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(),  
            nn.MaxPool2d(kernel_size=2),                                                
        )
  
        self.classifier = nn.Sequential(
            nn.Linear(in_features = 25088, out_features = 4096),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout(),            
            nn.Linear(in_features=4096, out_features=n_classes)
        )


    def forward(self, x):
        x = self.feature_extractor(x)
        x = torch.flatten(x, 1) # 1차원으로 쫙펴준다
        logits = self.classifier(x) # classifier에 통과시켜준다
        probs = F.softmax(logits, dim=1) # 나온 10개의 값들을 softmax로 확률 구해준다.
        return logits, probs

 

 

이제 어느 정도 준비는 다 끝났고, 본격적으로 이미지를 불러와 전처리 하고 학습을 진행해보자.

9. 데이터 불러오고 확인하기

# specify a data path
path2data = '/content/VGG/MyDrive/data'

# if not exists the path, make the directory
if not os.path.exists(path2data):
    os.mkdir(path2data)

# load dataset
train_ds = datasets.STL10(path2data, split='train', download=True, transform=transforms.ToTensor())
val_ds = datasets.STL10(path2data, split='test', download=True, transform=transforms.ToTensor())

plt.imshow(train_ds.data[1].T , cmap = 'gray_r') #대충 어떤 이미지가 있나 봅시다.
print(len(train_ds)) # train data의 개수
print(len(val_ds)) # test data의 개수
train_ds[0][0].shape

 

5000개의 training 데이터와 8000개의 validation 데이터가 있음을 확인할 수 있다. 이미지의 사이즈는 96*96 이다.

 

10. 이미지 전처리

먼저 효과적인 학습을 위해 데이터에 torchvision.transforms.Nomalize()를 진행하기 전에 필요한 parameter들을 얻자.

# To normalize the dataset, calculate the mean and std
train_meanRGB = [np.mean(x.numpy(), axis=(1,2)) for x, _ in train_ds]
train_stdRGB = [np.std(x.numpy(), axis=(1,2)) for x, _ in train_ds]

train_meanR = np.mean([m[0] for m in train_meanRGB])
train_meanG = np.mean([m[1] for m in train_meanRGB])
train_meanB = np.mean([m[2] for m in train_meanRGB])
train_stdR = np.mean([s[0] for s in train_stdRGB])
train_stdG = np.mean([s[1] for s in train_stdRGB])
train_stdB = np.mean([s[2] for s in train_stdRGB])


val_meanRGB = [np.mean(x.numpy(), axis=(1,2)) for x, _ in val_ds]
val_stdRGB = [np.std(x.numpy(), axis=(1,2)) for x, _ in val_ds]

val_meanR = np.mean([m[0] for m in val_meanRGB])
val_meanG = np.mean([m[1] for m in val_meanRGB])
val_meanB = np.mean([m[2] for m in val_meanRGB])

val_stdR = np.mean([s[0] for s in val_stdRGB])
val_stdG = np.mean([s[1] for s in val_stdRGB])
val_stdB = np.mean([s[2] for s in val_stdRGB])

print(train_meanR, train_meanG, train_meanB)
print(val_meanR, val_meanG, val_meanB)

#0.4467106 0.43980986 0.40664646
#0.44723064 0.4396425 0.40495726

 

그 뒤 transforms를 정의하고, 이미지를 여러장 불러와 다시 한번 확인해보자.

# transforms 정의하기
transforms_set = transforms.Compose([transforms.Resize((256,256)),transforms.RandomCrop((224,224)),transforms.ToTensor(),
                                  transforms.Normalize([train_meanR, train_meanG, train_meanB], [train_stdR, train_stdG, train_stdB])])



# 불러온 MNIS data 확인하기
ROW_IMG = 10
N_ROWS = 5

fig = plt.figure()
for index in range(1, ROW_IMG * N_ROWS + 1):
    plt.subplot(N_ROWS, ROW_IMG, index)
    plt.axis('off')
    plt.imshow(train_ds.data[index].T, cmap='gray_r')
fig.suptitle('For VGG Dataset - preview');

 

 

정의한 transforms를 train, validate 데이터셋에 적용하고, 변환된 이미지 사이즈를 확인하자.

train_ds.transform = transforms_set
val_ds.transform = transforms_set
train_ds[0][0].size()
#torch.Size([3, 224, 224])

 

이제 DataLoader만 정의하면 준비 끝이다!

# data loader 정의하기
train_loader = DataLoader(dataset=train_ds,
                          batch_size=Batch_size,
                          shuffle=True)

valid_loader = DataLoader(dataset=val_ds,
                          batch_size=Batch_size,
                          shuffle=False)

 

11. 학습

torch.manual_seed(Random_seed)

model = VGG16(N_classes).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001,weight_decay=0)
criterion = nn.CrossEntropyLoss()

 

이 부분에서 문제가 조금 있었다. 논문에서 제시한 대로 parameter을 설정했더니 학습이 진행되지 않았다. 그래서 learning rate과 weight decay 값을 위 코드와 같이 설정하였다. 또한 optimizer로 SGD 대신 Adam을 사용하였다.

 

model, optimizer, _ = training_loop(model, criterion, optimizer, train_loader,
                                    valid_loader, N_epochs, DEVICE)

 

이제 기다림의 시간만이 남아있다.

카페에서 돌리다가 너무 오래걸려서 집에 가서 다시 돌렸다 ㅠㅠ

 

 

에포크 20~50번째 까지의 학습 상황이다. Train Accuracy는 증가하고 있지만 Valid Accuracy는 fluctuate하고 있는 상황이었다. 오버피팅이 의심된다.

 

오버피팅이 일어난 것이 맞는 것 같다. 내 생각에는 15에포크 정도에서 Early stopping을 했어야 했던 것 같다. Keras에서는 Early stopping을 지원한다고 하니, 이를 이용하여 다시 한번 코드를 돌려보아야 할 것 같다. 

 

추가적으로, 이번에는 VGG16에 대해서만 실험을 진행하였는데, 다음 시간에는 논문에서 소개된 모든 모델에 대해서 간편하게 코딩할 수 있는 방법에 대해 소개하도록 하겠다.

 

끝!