[Pytorch] MNIST 데이터 셋 4가지 optimizer 성능 비교(+SGD, Adagrad, RMSprop, Adam)

지난번에 SGD 옵티마이저를 사용하여 MNIST 데이터 셋을 학습시키는 과정에 대해 살펴보았습니다. 이번에는 SGD를 포함해서 Adagrad, RMSprop, Adam 옵티마이저를 사용하여 MNIST 데이터 셋 학습을 진행해보려고 합니다. 코드 중간중간에 주석이 있기는 하지만 지난번 포스팅과 차이가 있는 부분만 설명드리려고 하니 조금 더 자세한 내용이 궁금하시다면 화면 하단의 SGD로 MNIST 데이터 셋 학습 링크를 참고해주시기 바랍니다. 또한 화면 하단 github 주소에 전체 코드를 올려두겠습니다. 

 

1. 모듈 및 라이브러리 import

import torch.nn as nn
import torch
import matplotlib.pyplot as plt
import numpy as np
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import torch.optim as optim
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

관련 모듈, 라이브러리 import를 진행합니다. 마지막에 os.environ[ ] 부분은 plt를 사용해서 그래프를 출력하려고 할 때 에러가 나면서 커널이 죽어 추가해준 것입니다.

 

2. Model 정의

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784,100)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(100,100)
        self.fc3 = nn.Linear(100,10)
    
    def forward(self, x):
        x1 = self.fc1(x)
        x2 = self.relu(x1)
        x3 = self.fc2(x2)
        x4 = self.relu(x3)
        x5 = self.fc3(x4)
        
        return x5

fully-connected model을 정의하였습니다. 총 4개의 층으로 되어있습니다. input layer, output layer를 제외하고 두 개의  hidden layer로 구성했습니다. Input layer에는 784차원의 데이터가 들어오고 output layer에서는 최종 10개의 데이터를 반환합니다.

 

3. MNIST 데이터 셋 다운로드

download_root = './MNIST_data'

dataset1 = datasets.MNIST(root = download_root,
                         train=True,
                         transform = transforms.ToTensor(),
                         download=True)

dataset2 = datasets.MNIST(root=download_root,
                         train=False,
                         transform = transforms.ToTensor(),
                         download=True)

MNIST 데이터 셋을 다운로드 받는 과정입니다. dataset1을 train 데이터로 사용하고 dataset2를 test 데이터로 사용하겠습니다. 전달되는 인자에 대한 설명은 이전 포스팅에 했기 때문에 궁금하신 분들은 화면 하단의 포스팅을 참고해주시기 바랍니다. 

 

4. Batch size 정의

batch_s = 100
# dataset1_loader의 len은 600
# dataset2_loader의 len은 100
dataset1_loader = DataLoader(dataset1, batch_size = batch_s)
dataset2_loader = DataLoader(dataset2, batch_size = batch_s)

batch size를 100으로 해주었습니다. dataset1의 크기는 6만(train), dataset2의 크기는 1만(test)입니다. 배치 사이즈를 100으로 해줌에 따라 각각의 dataset1_loader의 크기는 600(60000/100), dataset2_loader의 크기는 100(10000/100)이 됩니다.

 

5. 4가지의 optimizer 케이스 및 model 생성

model_dict = {}
loss_dict = {}
accuracy_dict = {}
# optimizer에 따른 학습 정도를 살펴볼 4가지 테스트 케이스
optimizer_case = ['SGD','Adam','AdaGrad','RMSprop']
for key in optimizer_case:
    model_dict[key] = Net()
    loss_dict[key] = []
    accuracy_dict[key] = []

각각의 optimizer 케이스마다 값을 담아줄 딕셔너리 3가지를 선언했습니다. optimizer의 이름을 key로 model, loss, accuracy를 value로 담아줄 것입니다. 옵티마이저는 'SGD', 'Adam', 'AdaGrad', 'RMSprop' 이렇게 네 가지를 적용해보려고 합니다. loss와 accuracy의 경우 나중에 학습 과정에서 여러 개의 값을 받을 것이기 때문에 value 부분을 list로 해주었습니다.

 

6. optimizer 정의

# 4가지 테스트케이스에 대한 optimizer 정의
optimizer_dict = {}
optimizer_dict['SGD'] = optim.SGD(model_dict['SGD'].parameters(),lr = 0.001 )
optimizer_dict['Adam'] = optim.Adam(model_dict['Adam'].parameters(),lr= 0.001)
optimizer_dict['AdaGrad'] = optim.Adagrad(model_dict['AdaGrad'].parameters(), lr=0.001)
optimizer_dict['RMSprop'] = optim.RMSprop(model_dict['RMSprop'].parameters(),lr=0.001)

각각의 옵티마이저를 정의하였습니다. learning rate는 모두 통일해서 0.001로 지정해주었고, 모멘텀은 따로 설정하지 않았습니다.

 

7. loss function 및 epoch 정의

# loss_function, total batch size, epoch 정의
loss_function = nn.CrossEntropyLoss()
total_batch = len(dataset1_loader) # 600 (60000 / 100) => (train dataset / batch_size)
epochs = np.arange(1,16)

loss_function은 CrossEntropy(크로스엔트로피)로 epoch은 총 15번으로 지정해주었습니다. 학습 과정에서 loss 값을 600번 더한 것을 전체 배치 크기로 나누어주기 위해 total_batch도 선언해주었습니다.

 

8. MNIST 데이터 셋 학습 및 loss, accuracy 구하기

# 총 4가지 optimizer를 사용하여 학습
for optimizer_name, optimizer in optimizer_dict.items():
    print(optimizer_name)
    for epoch in epochs:
        cost=0
        for images, labels in dataset1_loader: #dataloader는 image와 label로 구성
            # 하나의 Tensor에 데이터 784(28x28)개가 담긴 리스트가 100개 들어있음
            # 그리고 그것이 총 dataset1_loader의 len인 600개 존재
            images = images.reshape(100,784) 
           
            model_dict[optimizer_name].zero_grad()
            
            # feed forward
            predict = model_dict[optimizer_name].forward(images)
            
            # loss 값 구하기
            loss = loss_function(predict,labels) # 예측된 것과 label이 얼마나 차이가 나는지
            
            # back propagation
            loss.backward()
            
            # optimizer update
            optimizer.step()
            
            cost += loss # 총 600번의 loss를 더한다.
            
        with torch.no_grad(): # 미분하지 않겠다
            total = 0
            correct = 0
            for images, labels in dataset2_loader:
                images = images.reshape(100,784)
                
                outputs = model_dict[optimizer_name].forward(images)
                
                # torch.max에서 두 번째 인자는 dim을 의미
                # 1로 지정했다는 것은 하나의 행에서 가장 큰 값을 찾겠다는 것
                # dim을 지정하면 인덱스에 있는 값과 인덱스를 반환
                _,predict = torch.max(outputs, 1) 
                
                total += labels.size(0)
                correct += (predict == labels).sum() # 예측한 것과 labels이 얼마나 일치하는지
                
            avg_cost = cost / total_batch # loss 값 600개의 평균
            accuracy = 100 * (correct/total) 
            
            loss_dict[optimizer_name].append(avg_cost.detach().numpy())
            accuracy_dict[optimizer_name].append(accuracy)
            
            print("epoch : {} | loss : {:.6f}" .format(epoch, avg_cost))
            print("Accuracy : {:.2f}".format(100*correct/total))

- 학습 진행

총 4가지의 옵티마이저 경우를 따져주어야 하기 때문에 가장 바깥쪽에서 optimizer_dict만큼 반복을 진행합니다. dataset1_loader는 image와 label로 구성되어 있습니다. image가 input layer에 들어갈 수 있도록 차원을 맞춰주어야 합니다. reshape의 내용은 784개의 데이터가 들어 있는 1차원 백터가 총 100행 있는 것을 의미합니다. feed forward, loss 구하기, back propagation, optimizer update를 순서대로 진행한 뒤에 결과로 나온 loss를 더해줍니다.

 

- 테스트 진행

학습이 완료되었다면 완료된 모델에 test dataset을 넣어서 제대로 학습이 되었는지 확인합니다. 어떤 과정으로 테스트하는 지 궁금하신 분들은 주석 부분을 참고해주시기 바랍니다. 참고로 with torch.no_grad()부터 테스트 과정입니다. 최종 loss와 accuracy 결과를 loss_dict와 accuracy_dict에 담습니다.

 

9. 그래프 그리기

markers = {'SGD' : 'o', 'Adam' : 'x','AdaGrad' : 's', 'RMSprop' : 'D' }
plt.figure(figsize = (10,5))
plt.subplot(1,2,1)

for key in optimizer_case:
    plt.plot(epochs,loss_dict[key], marker = markers[key], markevery=100, label = key)
plt.xlabel("Epoch")
plt.ylabel("Loss")

plt.legend()

plt.subplot(1,2,2)
for key in optimizer_case:
    plt.plot(epochs, accuracy_dict[key],marker = markers[key], markevery=100, label=key)
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()

plt.show()

최종적으로 학습이 진행됨에 따라 loss와 accuracy 변화를 그래프로 그리는 과정입니다. 총 4개를 비교하기 때문에 marker를 사용하여 알아보기 쉽게 하였습니다.

 

그래프가 그려진 것을 보면 Adam이나 RMSprop가 AdaGrad나 SGD보다 좋은 성능을 보이는 것을 확인할 수 있습니다. 하지만 하이퍼 파라미터와 신경망의 구조에 따라서 optimizer 별로 성능 차이가 존재합니다. 위의 결과는 제가 정의한 하이퍼 파라미터와 신경망 구조에 대한 결과인 것입니다. 각 코드들에 대해 조금 더 자세한 설명을 원하시면 아래의 링크를 참고하시기 바랍니다. 

 

10. 참고하면 좋은 글

[ 자세한 설명 ] SGD로 MNIST 데이터 셋 학습(+loss 및 accuracy)

[Pytorch] jupyter notebook으로 MNIST 데이터 셋 학습(+정확도, loss 측정)

[실습 전체 코드] Github 사이트

LHI0915/Pytorch_tutorials