[모각코] MNIST 분류 모델 리뷰

저번 시간엔 DNN에 대해 배워보았다. 이번시간엔 MNIST 데이터셋을 분석 하는 인공지능 모델을 리뷰해보는 시간을 가져보자.

이렇게 총 4개의 파일로 이루어져있다. 먼저 models.py 부터 살펴보자.

Data_loader.py

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def load_data_loader(config):
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = datasets.MNIST(config.data_dir, train=True, download=True, transform=transform)
    valid_dataset = datasets.MNIST(config.data_dir, train=False, download=True, transform=transform)

    train_loader = DataLoader(dataset=train_dataset,batch_size=config.batch_size, shuffle=True)
    valid_loader = DataLoader(dataset=valid_dataset,batch_size=config.batch_size,shuffle=False)

    return train_loader, valid_loader

MNIST 데이터셋을 불러오는 파일이다

 

Models.py

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

class DNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, 64)
        self.fc5 = nn.Linear(64, 32)
        self.fc6 = nn.Linear(32, 10)

    def forward(self, x):
        x = x.float()
        h1 = F.relu(self.fc1(x.view(-1, 784)))
        h2 = F.relu(self.fc2(h1))
        h3 = F.relu(self.fc3(h2))
        h4 = F.relu(self.fc4(h3))
        h5 = F.relu(self.fc5(h4))
        h6 = self.fc6(h5)
        return F.log_softmax(h6, dim=1)

28*28 픽셀사이즈의 손글씨데이터를 1*784 사이즈로 변환시켜 fc1을통해 사이즈 512로 변환한뒤 활성화함수를 거친다. 이과정을 반복하여 사이즈를 10까지 변환시켜준다. 이 사이즈 10짜리의 데이터는 softmax 함수를 통해 각각 0부터 9까지의 확률을 나타내게된다.

 

Train.py

import argparse 
import torch 
import torch.nn as nn
import torch.optim as optim
from models import DNN
from data_loader import load_data_loader
from trainer import MnistTrainer

def define_argparser():
    p = argparse.ArgumentParser() 
    p.add_argument('--data_dir', type=str,default='./Mnist')    
    p.add_argument('--batch_size',type=int,default=512)
    p.add_argument('--n_epochs',type=int,default=10)
    p.add_argument('--lr',type=float,default=0.01)
    p.add_argument('--seed',type=int,default=42)
    
    config = p.parse_args()

    return config

def seed_everything(seed: int):
    import random, os
    import numpy as np
    import torch
    
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

def main(config):
    seed_everything(config.seed)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    train_loader, valid_loader = load_data_loader(config)

    print("Train:", len(train_loader.dataset))
    print("Valid:", len(valid_loader.dataset))

    model = DNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=config.lr)
    crit = nn.CrossEntropyLoss()

    trainer = MnistTrainer(config, model, crit, optimizer, train_loader, valid_loader, device)
    trainer.train()

if __name__ == '__main__':
    config = define_argparser()
    main(config)

본격적으로 학습하기에 앞서 설정들을 하는 단계 seed는 랜덤한 값들을 고정시켜주는 작업이다. 벤치마킹을 제대로 하려면 같은 어느 환경에서든 같은 값들을 사용해야 하기 때문이라고 한다. main 에선 데이터셋을 불러오고 model, optimizer, loss_fn 등을 정해준뒤 trainer를 돌린다.

Trainer.py

import torch
import torch.nn.functional as F

class MnistTrainer():
    def __init__(self, config, model, crit, optimizer, train_loader, valid_loader, device):
        self.config = config
        self.model = model
        self.crit = crit
        self.optimizer = optimizer
        self.train_loader = train_loader
        self.valid_loader = valid_loader
        self.device = device
    
    def run_train(self, epoch):
        self.model.train()
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = F.nll_loss(output, target)
            loss.backward()
            self.optimizer.step()

            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(self.train_loader.dataset),
                100. * batch_idx / len(self.train_loader), loss.item()))

    def run_valid(self):
        self.model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in self.valid_loader:
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)
                test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
                pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
                correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(self.valid_loader.dataset)

        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(self.valid_loader.dataset),
            100. * correct / len(self.valid_loader.dataset)))
                
    def train(self):
        for epoch in range(1, self.config.n_epochs+1):
            self.run_train(epoch)
            self.run_valid()

run_train을 통해 학습을 한뒤, 완료 후 run_valid를 통해 전체 손실정도와 정확성을 계산하여 출력해준다..