PyTorch-실습 : 딥러닝 모델 프로그램 만들기(2)

이전 포스팅을 보시면 아시겠지만, 딥러닝을 구현 단계는 크게 3단계로 나눌 수 있습니다.

1. 데이터, 2 모델, 3 학습

과정을 정리하면,

 

데이터를 가져오고,

그에 맞는 모델을 만들어서

학습(훈련)을 하게 됩니다.

 

그렇기 때문에, 하나의 딥러닝 프로젝트를 진행한다면 3개의 파일(data.py, model.py, trainer.py)로 나누어서, 코드를 작성하시고, 이러한 3개 코드를 컨트롤할 수 있는, 메인 코드(train.py)를 만드시는것을 추천 드립니다.

 

그동안 한 파일에 모든 코드를 작성을 하여 프로젝트를 진행하셨을 겁니다. 그렇게 되다 보니 활용이나 응용 과정에서 어려움을 많이 겪게 되고 문제가 생겼을 시 디버깅이 어려운 문제를 갖게 될것입니다.

 

실제 업무에서나 프로젝트 진행시, 이렇게 파일을 나누고 필요에 따라 더 세분화합니다. 이렇게 나누어서 학습을 진행하셔야 나중에 내가 데이터만 바꾸어서 다른 프로젝트도 응용이 가능하고, 필요에 따라서 모델만 가지고 다른 프로젝트에 활용이 가능합니다.

 

 

data.py

0.00MB

model.py

0.00MB

train.py

0.00MB

trainer.py

0.00MB

model.py

import torch.nn as nn

class MnistModel(nn.Module):

    def __init__(self, input_size, output_size):
        self.input_size = input_size
        self.output_size = output_size
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, 500), nn.LeakyReLU(), nn.BatchNorm1d(500),
            nn.Linear(500, 300), nn.LeakyReLU(), nn.BatchNorm1d(300),
            nn.Linear(300, 100), nn.LeakyReLU(), nn.BatchNorm1d(100),
            nn.Linear(100, 50), nn.LeakyReLU(), nn.BatchNorm1d(50),
            nn.Linear(50, output_size),
            nn.LogSoftmax(dim=-1))

    def forward(self, x):
        y = self.layers(x)
        return y

 

data.py

from torchvision import datasets, transforms

def load_mnist():
    dataset = datasets.MNIST('../data', train=True, download=True, transform=transforms.Compose(transforms.ToTensor()))
    x_data = dataset.data.float() / 255.
    y_data = dataset.targets
    return x_data, y_data

 

trainer.py

from copy import deepcopy
import numpy as np
import torch


class Train_model():
    def __init__(self, model, optimizer, loss):
        self.model = model
        self.optimizer = optimizer
        self.loss = loss
        super().__init__()

    def run_train_data(self, x, y, config):
        self.model.train()
        indices = torch.randperm(x.size(0))
        x = torch.index_select(x, dim=0, index=indices).split(config.batch_size, dim=0)
        y = torch.index_select(y, dim=0, index=indices).split(config.batch_size, dim=0)
        train_loader = zip(x, y)
        total_loss = 0
        for i, (data, target) in enumerate(train_loader):
            train_model = self.model(data)
            train_loss = self.loss(train_model, target.squeeze())
            self.optimizer.zero_grad()
            train_loss.backward()
            self.optimizer.step()
            print(f"Train_{i} / loss : {float(train_loss)}")
            total_loss += float(train_loss)
        return total_loss / len(x)

    def run_test_data(self, x, y, config):
        self.model.eval()
        # no_grad mode
        with torch.no_grad():
            indices = torch.randperm(x.size(0))
            x = torch.index_select(x, dim=0, index=indices).split(config.batch_size, dim=0)
            y = torch.index_select(y, dim=0, index=indices).split(config.batch_size, dim=0)
            total_loss = 0
            test_loader = zip(x, y)
            for i, (data, target) in enumerate(test_loader):
                test_model = self.model(data)
                test_loss = self.loss(test_model, target.squeeze())
                print(f"test_{i} / loss : {float(test_loss)}")
                total_loss += float(test_loss)
            return total_loss / len(x)

    def main_train(self, train_data, test_data, config):
        result_loss = np.inf
        best_model = None
        for i in range(config.epochs):
            train_loss = self.run_train_data(train_data[0], train_data[1], config)
            test_loss = self.run_test_data(test_data[0], test_data[1], config)
            if test_loss <= result_loss:
                result_loss = test_loss
                best_model = deepcopy(self.model.state_dict())
            print(f"Epoch : {i+1}, train_loss : {train_loss}, test_loss : {test_loss}  result_loss : {result_loss}")
        self.model.load_state_dict(best_model)

 

train.py

import torch
import argparse
import torch.nn as nn
import torch.optim as optim
from model import MnistModel
from trainer import Train_model
from data import load_mnist


def user_option():
    page = argparse.ArgumentParser()
    page.add_argument('--save_model', required=True)
    page.add_argument('--train_ratio', type=float, default=.8)
    page.add_argument('--batch_size', type=int, default=64)
    page.add_argument('--epochs', type=int, default=5)
    config = page.parse_args()
    return config


def mnist_run(config):
    device = torch.device('cpu')
    x_data, y_data = load_mnist()
    x_data = x_data.view(x_data.size(0), -1)
    train_cnt = int(x_data.size(0) * config.train_ratio)
    test_cnt = x_data.size(0) - train_cnt
    indices = torch.randperm(x_data.size(0))
    x_data = torch.index_select(x_data,dim=0,index=indices).to(device).split([train_cnt, test_cnt], dim=0)
    y_data = torch.index_select(y_data,dim=0,index=indices).to(device).split([train_cnt, test_cnt], dim=0)
    model = MnistModel(784, 10).to(device)
    optimizer = optim.Adam(model.parameters())
    loss = nn.NLLLoss()
    trainer = Train_model(model, optimizer, loss)
    trainer.main_train((x_data[0], y_data[0]), (x_data[1], y_data[1]), config)
    torch.save({'model': trainer.model.state_dict(), 'config': config}, config.save_model)


if __name__ == '__main__':
    config = user_option()
    mnist_run(config)