[pytorch] VGGNet으로 ImageNet 학습하기

1. VGGNet이란?

 

CNN 알고리즘들 중에서는 이미지 분류(Image Classification)용 알고리즘, 예를 들어 AlexNet, GoogleNet 등이 있다.

그 이미지 분류 CNN 모델들 중에 하나가 바로 VGGNet이다.

 

VGGNet은 몇 개의 층(layer)으로 구성되어 있는지에 따라, 16개 층으로 구성되어 있으면 VGG16,

19개 층으로 구성되어 있으면 VGG19라고 불린다.

 

VGGNet은 옥스포드 대학의 연구팀 VGG에 의해 개발된 모델로써,

2014년 이미지넷 이미지 인식 대회에서 준우승을 한 모델이다.

 

2012년, 2013년 우승 모델들은 8개의 층으로 구성되어 있는 반면,

2014년 VGGNet(VGG19)는 19층으로 구성되었고, GoogleNet은 22층으로 구성되었다.

그리고 2015년에 이르러서는 152개의 층으로 구성된 ResNet이 제안되었다.

따라서 VGGNet 모델부터 시작해서 네트워크의 깊이가 확 깊어진 것을 확인할 수 있다.

 

네트워크의 깊이 변화에 따른 top-5 error의 변화

 

VGGNet은 사용하기 쉬운 구조와 좋은 성능 덕분에 그 대회에서 우승을 거둔 조금 더 복잡한 형태의 GoogleNet보다 더 인기를 끌었다.

 

VGGNet의 original 논문의 링크는 아래에 걸어두겠다.

: https://arxiv.org/pdf/1409.1556.pdf

불러오는 중입니다...

 

2. VGGNet의 구조

 

VGGNet 연구의 핵심은 "네트워크의 깊이"가 성능에 어떤 영향을 미치는지를 확인하고자 한 것이다.

따라서 VGG 연구팀은 깊이의 영향만을 확인하고자, 컨볼루션 필터 커널의 사이즈를 가장 작은 3x3으로 고정했다.

 

VGGNet 구조 설명 표 [출처: original 논문]

VGG 연구팀은 총 6개의 구조(A, A-LRN, B, C, D, E)를 만들어 성능을 비교했다.

이들 중, D구조가 VGG16이고 E구조가 VGG19라고 보면 된다.

 

이들은 깊이가 11층, 13층, 16층, 19층으로 깊어지면서 분류 에러가 감소하는 것을 관찰했다.

즉, 깊어질수록 성능이 좋아진다는 것이다.

 

다음은 VGG16의 구조를 그림으로 나타낸 것이다.

 

VGG16 구조

그림에서 확인할 수 있듯이, VGG는 합성곱 계층(검정색)과 풀링 계층(빨간색)으로 구성된다.

다만, 합성곱 계층(검정색)과 완전연결 계층(파란색)을 합쳐서 모두 16층(VGG19의 경우 19층)으로 심화한 것이 특징이다.

 

인풋으로는 224 x 224 x 3 이미지(224 x 224 RGB 이미지)를 입력받을 수 있다.

 

인풋 값이 16개의 층을 지난 후 softmax 함수로 활성화 된 출력값들은 1000개의 뉴런으로 구성된다.

이 말은 즉, 1000개의 클래스로 분류하는 목적으로 만들어진 네트워크라는 뜻이다.

 

 

3. VGGNet으로 ImageNet 학습하기

3-1. Data 준비하기

 

일단 학습할 ImageNet dataset을 다운로드 받아야 한다.

 

http://image-net.org/download <- 여기에서 2012년도 dataset을 다운로드 받는다.

data의 크기가 워낙 크기 때문에 다운로드 받는데 3~5일이 걸렸다..

ImageNet

ImageNet Dataset 파일

 

 

이 데이터를 사용할 수 있게 data를 load해줘야 한다.

import torch
import torchvision
import torchvision.transforms as transforms

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

import matplotlib.pyplot as plt
import numpy as np

import json
import os

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

trainset = torchvision.datasets.ImageNet('../../data/imagenet', split='train', download=None, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)

testset = torchvision.datasets.ImageNet('../../data/imagenet', split='val', download=None, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

이렇게 trainset과 testset을 나눠서 torchvision.datasets.ImageNet으로 data를 load해준다.

 

(pytorch에서 지원하고 있는 dataset의 종류마다 torchvision.datasets 사용법이 다르니,

pytorch 공식 사이트인 https://pytorch.org/docs/stable/torchvision/datasets.html#imagenet 를 참고한다.)

 

위 코드를 실행시키면 아래와 같이 데이터셋이 train 폴더와 val 폴더로 분리되어 생성된다.

각 폴더 안에는 이렇게 데이터셋의 각 class가 폴더로 구성되어 있고,

각 class 폴더 안에는 그 class의 image 파일들이 있다.

위 사진은 class index가 "tench"(물고기 종류 중 하나)인 이미지들이다.

 

이제 학습시킬 data 준비가 완료되었다.

 

 

3-2. VGG 모델 구현하기

 

나는 vgg19 모델을 사용하였다.

 vgg19의 코드는 다음과 같다.

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Sequential(
            #3 224 128
            nn.Conv2d(3, 64, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(64, 64, 3, padding=1),nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            #64 112 64
            nn.Conv2d(64, 128, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(128, 128, 3, padding=1),nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            #128 56 32
            nn.Conv2d(128, 256, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(256, 256, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(256, 256, 3, padding=1),nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            #256 28 16
            nn.Conv2d(256, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(512, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(512, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2),
            #512 14 8
            nn.Conv2d(512, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(512, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.Conv2d(512, 512, 3, padding=1),nn.LeakyReLU(0.2),
            nn.MaxPool2d(2, 2)
        )
        #512 7 4

        self.avg_pool = nn.AvgPool2d(7)
        #512 1 1
        self.classifier = nn.Linear(512, 1000)
        """
        self.fc1 = nn.Linear(512*2*2,4096)
        self.fc2 = nn.Linear(4096,4096)
        self.fc3 = nn.Linear(4096,10)
        """

    def forward(self, x):
        #print(x.size())
        features = self.conv(x)
        #print(features.size())
        x = self.avg_pool(features)
        #print(avg_pool.size())
        x = x.view(features.size(0), -1)
        #print(flatten.size())
        x = self.classifier(x)
        #x = self.softmax(x)
        return x, features

간단히 쓱 훑어보아도 layer가 19개로 구성되어 있다는 것을 확인할 수 있다.

 

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = Net()
net = net.to(device)
param = list(net.parameters())

그리고 위 코드의 첫번째 줄을 통해 device라는 변수명에 GPU 디바이스를 지정/저장해준다.

torch.cuda.is_available() 은 GPU를 사용할 수 있는지 반환하는 메소드이다.

그다음, net=Net() 을 통해 Net을 초기화해주고,

net.to(device) 구문을 통해 해당 네트워크를 GPU에서 돌아가는 네트워크로 지정한다.

 

import torch.optim as optim

criterion = nn.CrossEntropyLoss().cuda()
optimizer = optim.Adam(net.parameters(),lr=0.00001)

idx2label = []
cls2label = {}
with open("./imagenet_class_index.json", "r") as read_file:
    class_idx = json.load(read_file)
    idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))]
    cls2label = {class_idx[str(k)][0]: class_idx[str(k)][1] for k in range(len(class_idx))}

그 후 criterion과 optimizer를 정의한다.

criterion은 일반적으로 loss라고 부르며, 이를 통해 학습을 위한 loss값을 계산하게 된다.

 

Adam optimzer는 현재 자주 사용되는 옵티마이저로,

학습 시 현재의 미분값 뿐만이 아니라, 이전 결과에 따른 관성 모멘트를 가지고 있는 것이 특징이다.

 

그다음, imagenet_class_index.json 파일을 다운로드 받아 열어준다.

ImageNet은 총 1000개의 class로 구성되어 있고 단순히 "cat"  "dog" 가 아니라 "n01440764" 와 같은

직관적으로 파악하기 어려운 index로 구분되어 있기 때문에 이를 알기 쉬운 label로 변환하는 작업이다.

 

다음은 imagenet_class_index.json 파일의 구조이다.

이 파일은 구글링으로 쉽게 다운받을 수 있을 것이다.

 

imagenet_class_index.json

 

 

3-3. VGG 모델 실행하기

 

실행 코드는 다음과 같다.

for epoch in range(3):  # loop over the dataset multiple times
    running_loss = 0.0

    if(epoch>0):
        net = Net()
        net.load_state_dict(torch.load(save_path))
        net.to(device)

    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        # zero the parameter gradients
        optimizer.zero_grad()
        outputs,f = net(inputs)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        if(loss.item() > 1000):
            print(loss.item())
            for param in net.parameters():
                print(param.data)
        # print statistics
        running_loss += loss.item()
        if i % 50 == 49:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 50))
            running_loss = 0.0
    save_path="/data/data1/minjoo/workspace/vgg/imagenet_vgg16_result.pth"
    torch.save(net.state_dict(), save_path)
        #print("\n")

print('Finished Training')

class_correct = list(0. for i in range(1000))
class_total = list(0. for i in range(1000))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images = images.cuda()
        labels = labels.cuda()
        outputs,_ = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(16):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

accuracy_sum=0
for i in range(1000):
    temp = 100 * class_correct[i] / class_total[i]
    print('Accuracy of %5s : %2d %%' % (
        idx2label[i], temp))
    accuracy_sum+=temp
print('Accuracy average: ', accuracy_sum/1000)

나는 빠르게 실행시키기 위해 epoch을 일단 3으로 해줬는데 더 많이 실행시킬수록 loss가 줄어드니 좋을 것이다.

 

그리고 1 epoch이 실행될 때마다 net을 저장해주고 다음 epoch을 돌 때 그 net을 불러온다.

이는 중간에 실행을 끊더라도 다시 실행시킬 때 처음부터 training되는 것이 아닌, 끊은 시점부터 training되도록 하기 위함이다.

 

아래 사진은 2 epoch째 실행되고 있는 사진이다.

처음에는 loss가 6.9였는데 1정도 줄어든 것을 확인할 수 있다.

training이 완료되면 accuracy를 계산하여 각 class의 정확도를 print해주고,

추가로 1000개 class의 accuracy의 평균도 계산하여 마지막에 print해준다.

 

 

 

 

* 참고한 사이트

https://poddeeplearning.readthedocs.io/ko/latest/CNN/VGG19%20+%20GAP%20+%20CAM/

VGG Net 과 CAM 구현 - POD_Deep-Learning

https://tutorials.pytorch.kr/

파이토치(PyTorch) 튜토리얼에 오신 것을 환영합니다 — PyTorch Tutorials 1.4.0 documentation