MNIST 데이터 CNN으로 분류하기 - review

예전부터 계속 파이토치를 써오기는 했지만, 다시 공부하는 차원에서 CNN을 이용한 mnist 분류를 리뷰해보고자 한다.

 

1. 필요한 라이브러리 호출하기

import os
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms, datasets

torch -> pytorch 불러오기

numpy -> linear algebra를 계산 위한 넘파이

torch.nn -> layer, function들이 기본적으로 들어가있음

torch.optim -> 최적화 함수에 대한 기능들 내장

DataLoader -> 데이터 부를 때 사용함, train_loader, test_loader로 보통 만들어서 많이 사용

SummaryWriter -> 텐서보드에 그래프 그리기위함

torchvision -> 비전 데이터 관련해서 다룰 수 있는 툴들 및 데이터 셋이 많음

 

 

2. 하이퍼 파레미터 설정하기

lr = 1e-03
batch_size = 64
ckpt_dir = './checkpoint'
log_dir = './log'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

lr -> Learning rate

batch_size -> 학습할 때 몇 개의 데이터씩 넣을 지?

ckpt_dir -> checkpoint 저장할 경로

log_dir -> 텐서보드를 저장할 디렉토리

device -> torch cpu or gpu

 

 

3. Convolution Network 만들기

class Convnet(nn.Module):
	def __init__(self):
    	super(Convenet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size = 5, stride=1, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.relu1 = nn.ReLU()
        
        self.conv2 = nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5, stride=1, padding=0, bias=True)
        self.drop2 = nn.Dropout2d(0.5)
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        self.relu2 = nn.ReLU()
        
    	self.fc1 = nn.Linear(20*4*4, 50, bias = True)
        self.relu1_fc1 = nn.ReLU()
        self.drop1_fc1 = nn.Dropout2d(p=0.5)
        
        self.fc2 = nn.Linear(in_features=50, out_features=10, bias=True)
        
    def forward(self,x):
    	x = self.conv1(x)
        x = self.pool1(x)
        x = self.relu(x)
        
        x = self.conv2(x)
        x = self.drop2(x)
        x = self.pool2(x)
        x = self.relu2(x)
        
        x = x.view(-1,320)
        x = self.fc1(x)
        x = self.relu1_fc1(x)
        x = self.drop1_fc1(x)
        
        x = self.fc2(x)
        
        return x

intput -> Conv1 -> pooling -> Conv2 -> pooling -> FC1 -> FC2 -> output

input_shape - 1,28,28 (channel, height, width)

conv1, pooling 연산 후 , shape = (1, 12, 12)

conv2, pooling 연산 후 , shape = (1, 4, 4)

=> feature map 갯 수 20 * 4* 4 => 320

 

 

4. 만든 모델의 state 저장과 load를 위한 함수 생성

def save(ckpt_dir, net, optim, epoch):
	if not os.path.exists(ckpt_dir):
    	os.makedirs(ckpt_dir)
        
    torch.save({'net' : net.state_dict(), 'optim' : optim.state_dict()},
    			'./%s/model_epoch%d.pth' %(ckpt_dir, epoch))
                
def load(ckpt_dir, net, optim):
	ckpt_lst = os.listdir(ckpt_dir)
    ckpt_lst.sort()
    
    dict_model = torch.load('./%s/%s' %(ckpt_dir, ckpt_lst[-1]))
    
    net.load_state_dict(dict_model['net])
    optim.load_state_dict(dict_model['optim])
    
    return net, optim

 

5. 데이터 불러오기

transform = transforms.Compose([
	transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))
])

dataset = datasets.MNIST(download=True, root='./', train=True, transform=transform)
loader = DataLoader(dataset, batch_size = batch_size, shuffle=True, num_workers=0)

num_data = len(loader.dataset)
num_batch = np.ceil(num_data/batch_size)

6. 손실 및 최적화 함수 정의, 네트워크 훈련

net = Convnet().to(device)
params = net.parameters()

criterion = nn.CrossEntropyLoss().to(device)
fn_pred = lambda output:torch.softmax(output, dim=1)
fn_acc =lambda pred, label : ((pred.max(dim=1)[1]==label).type(torch.float)).mean()
optim = torch.optim.Adam(params, lr=lr)
writer = SummaryWriter(log_dir = log_dir)

for epoch in range(1, num_epoch + 1):
  net.train()

  loss_arr = []
  acc_arr = []

  for batch, (input, label) in enumerate(loader, 1):
    input = input.to(device)
    label = label.to(device)

    output = net(input)
    pred = fn_pred(output)

    optim.zero_grad()

    loss = criterion(output, label)
    acc = fn_acc(pred, label)
    loss.backward()

    optim.step()

    loss_arr += [loss.item()]
    acc_arr += [acc.item()]

    print(f'Train : epoch : {epoch}/{num_epoch} | batch : {batch}/{num_batch} | loss : {np.mean(loss_arr)} | acc : {np.mean(acc_arr)}')


  writer.add_scalar('loss', np.mean(loss_arr),epoch)
  writer.add_scalar('acc',np.mean(loss_arr), epoch)

  save(ckpt_dir = ckpt_dir, net=net, optim=optim, epoch=epoch)

writer.close()

 

 

7. Test data 불러오고 확인하기

-> 위의 dataloader에서 조금수정하고 훈련과정서도 훈련하는 process만 없애면 된다.

transform = transforms.Compose([
  transforms.ToTensor(),
  transforms.Normalize(mean=(0.5,), std=(0.5,))
])

dataset = datasets.MNIST(download=True, root='./', train=False, transform=transform)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)

num_data = len(loader.dataset)
num_batch = np.ceil(num_data/ batch_size)

net, optim = load(ckpt_dir=ckpt_dir, net= net,optim=optim)

with torch.no_grad():
  net.eval()

  loss_arr = []
  acc_arr = []

  for batch, (input, label) in enumerate(loader, 1):
    input = input.to(device)
    label = label.to(device)

    output = net(input)
    pred = fn_pred(output)

    #optim.zero_grad()

    loss = criterion(output, label)
    acc = fn_acc(pred, label)
    #loss.backward()

    optim.step()

    loss_arr += [loss.item()]
    acc_arr += [acc.item()]

    print(f'Test: batch : {batch}/{num_batch} | loss : {np.mean(loss_arr)} | acc : {np.mean(acc_arr)}')