cifar10 classification

이번에는 cifar10 데이터셋과 cnn을 이용하는 모델 분류를 리뷰해 보려고 한다.
이것도 매우 기본적인 cnn이지만 다시 복습하는데 의의를 두고 정리해 보겠다.

data_dir='/content/drive/MyDrive/machine_learning_project/cifar10/Dataset'
data_split_dir='/content/drive/MyDrive/machine_learning_project/cifar10/DataSplit/'

!pip install split_folders
import splitfolders

ratio_data = 0.7
splitfolders.ratio(data_dir, output=data_split_dir, seed=200, ratio=(ratio_data, 1-ratio_data))

data download 및 split

데이터를 미리 다운받았다는 전제하에 데이터의 경로를 설정하고
splitfolders를 사용하여 데이터를 train과 test를 7:3 비율로 나눈다.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torchvision
from torchvision import datasets, models, transforms

import time
import os
import copy
import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

data pre-processing

본격적으로 데이터와 모델을 위한 패키지를 가져온다.
데이터를 가져와서 증강시키고 load 하는 과정을 보자

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(256),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# train/val일때 각각 transform된 image_datasets 생성
image_datasets = {x: datasets.ImageFolder(os.path.join(data_split_dir, x), data_transforms[x]) for x in ['train', 'val']}

# 만들어진 Image_datasets의 dataloader 생성
dataloaders = {
    x: torch.utils.data.DataLoader(image_datasets[x], 
                                   batch_size=4, 
                                   shuffle=True, 
                                   num_workers=4) for x in ['train', 'val']
    }


dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
classes = ['Airplane','Bird','Frog']
class_names = image_datasets['train'].classes

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

① 먼저, 데이터 양을 늘리기 위해 data_transforms를 정의하고,
② data_transforms['기존 데이터'] -> image_datasets을 만든다.
③ 만들어진 데이터셋을 batch size = 4인 dataloader에 넣는다.
④ 데이터의 사이즈와 3가지(비행기, 새, 개구리) class를 정의한다.

• 데이터셋 구조 : image_datasets > train/test > classes > data

더 자세한 설명은 --> https://velog.io/@heyme/transforms-datasets-dataloader

model define

model = models.resnet18(pretrained=True)

# fc layer의 입력채널 수
num_ftrs = model.fc.in_features

# fc layer 수정
model.fc = nn.Linear(num_ftrs, len(class_names)) 
# nn.Linear(in_features, out_features)

model = model.to(device)

모델은 pretrain된 resnet을 사용한다. 여기서 우리가 원하는 모델 출력은
3개의 class에 관한 것이므로 resnet의 마지막 fc layer를 수정해야 한다.

먼저 model.fc.in_features을 이용하여 fc layer의 입력 채널 수를 구하고,
nn.Linear를 이용하여 fc layer의 입출력을 원하는 값으로 변경할 수 있다.

loss function & optimizer

criterion = nn.CrossEntropyLoss()

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

loss function은 분류이므로 CrossEntropyLoss, optimizer는 SGD를 사용하였다. scheduler는 stepLR을 사용하고 lr은 0.01에서 시작하여 7epoch마다 0.1의 비율로 줄어들게 하였다.

model train

이제 모델 학습을 위한 준비물은 다 준비되었고,
학습만 시키면 된다! 코드를 하나씩 뜯어보겠다!

train_losses = []
val_losses = []

def train_model(model, criterion, optimizer, scheduler, n_epochs=25):
  best_model_wts = copy.deepcopy(model.state_dict())
  best_acc = 0.0
  
  # 각 epoch마다 train/val 수행
  for epoch in range(n_epochs):
    for phase in ['train', 'val']:
      
      if phase == 'train':
        model.train()
      elif phase == 'val':
        model.eval()  

모델을 학습시킬 함수를 만든다. 필요한 준비물은 앞서 정의한
model, loss function, optimizer, lr_scheduler, epoch의 수 이다.

또한 train, val에 따라 train()과 eval()로 나누는 이유가 있는데, eval()은 그 과정중에 사용하면 안되는 layer를 자동으로 off 하도록 한다. 그런 layer로는 Dropout, BatchNorm 이 있다.
둘 다 과적합을 막고 학습이 잘 되게 돕는 regularization 역할이므로 학습이 아닌 평가에는 사용하지 않는다.

      running_loss = 0.0
      runing_corrects = 0

      for inputs, labels in dataloaders[phase]:
        inputs, labels = inputs.to(device), labels.to(device)

        with torch.set_grad_enabled(phase=='train'):
          outputs = model(inputs)
          _, preds = torch.max(outputs, dim=1)
          loss = criterion(outputs, labels)

          if phase == 'train':
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        running_loss += loss.item() * inputs.size(0)
        runing_corrects += torch.sum(preds == labels.data)

여기서부터 모델 학습에 관한 이야기이다. 틀린 개수와 맞은 개수의 변수를 선언하고 dataloader를 이용하여 한번에 데이터의 feature와 label을 4개씩, batch size 단위로 공급한다.

학습시 연산 기록을 추적하기 위해 torch.set_grad_enabled(phase=='train')을 사용하고 data를 모델의 입력으로 넣었을때 출력으로 예측값인 outputs을 받는다. outputs과 정답값인 labels을 손실함수인 crossentropy 입력으로 넣고 그 때의 loss를 구한다.

이제 loss를 알았으니 backpropagation을 진행한다. 먼저 그 전 loop에서 반영된 w의 기울기를 0으로 초기화하고, loss.backward()로 구한 ∂(loss)/∂w을 이용하여 w를 계산하고 업데이트 한다. backpropagation 설명 -> https://velog.io/@heyme/backpropagation
하나의 batch를 다 돌면, 각각 맞은 개수와 loss를 누적하여 계수한다.

      if phase == 'train':
        scheduler.step()

      # 해당 epoch의 정확도
      epoch_loss = running_loss / dataset_sizes[phase]
      epoch_acc = runing_corrects.double() / dataset_sizes[phase]    

      if phase == 'train':
        train_losses.append(loss.item())
      elif phase == 'val':
        val_losses.append(loss.item())
        
      if phase == 'val' and epoch_acc > best_acc:
        best_acc = epoch_acc
        best_model_wts = copy.deepcopy(model.state_dict())  
        
      print(f'[{phase}] [{epoch + 1} / {n_epochs}] Loss : {loss.item()}')
      
  print(f'Best accuracy is {best_acc}')
  model.load_state_dict(best_model_wts)
  return model

한 epoch이 끝나면, lr 값을 업데이트한다. 그리고 batch마다 계산한 loss와 맞은 개수를 전체 데이터 크기로 나눠 epoch 당 loss와 정확도의 평균을 구한다.

전체 epoch까지 끝나면 가장 정확도가 높은 모델의 가중치, best_model_wts을 적용한 모델을 출력으로 얻는다. 이렇게 train_model 함수의 구현이 끝났다. 이 함수를 돌려 모델을 학습한 결과는 아래와 같다.

train_model(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler)

plt.figure(figsize=(12,3))
plt.plot(val_losses)

이제 학습된 모델로 예측을 해보려고 한다. 예측 결과를 시각화하여 보여주기 위해
tensor로 변환되어 있는 데이터를 다시 이미지로 되돌리는 tensor_to_img를 사용한다.

def tensor_to_img(inp):
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1) * 255
    inp = inp.astype(np.uint8)
    return inp

그리고 예측한 결과를 보여주는 함수를 보면

def imshow_(inp, title=None):
    inp = tensor_to_img(inp)
    if title is not None:
        inp = np.pad(inp, ((40,0),(0,0), (0,0)), mode='constant', constant_values=0).copy() 
        inp = cv2.putText(inp, f'{title}', (0,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1)  
    inp = Image.fromarray(inp)
    return inp

def imshow(inp, title=None, show=True):
    inp = imshow_(inp, title) 
    if show:
        display(inp)
    else:
        return inp
        
def visualize_model(model, num_images=4):
    was_training = model.training
    # 모델평가
    model.eval()

    imgs = []
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            # image label과 예측 label이 포함된 img 시각화
            for inp, label, pred in zip(inputs, labels, preds):

                label = class_names[label]
                pred = class_names[pred]
                img  = imshow(inp.cpu(), f'{label}, {pred}', show=False)
                imgs.append(img)
            
            if num_images <= len(imgs):
                break

    ipyplot.plot_images(imgs[:num_images])

임의로 하나의 batch만 돌린 결과, outputs의 결과로 아래와 같은 tensor를 얻는다.
tensor([[10.5091, -5.6069, -5.9830],
[-6.2081, 6.0555, -0.7366],
[ 1.3087, -3.7720, 2.9738],
[-1.8460, -0.8243, 3.5677]], device='cuda:0')

이때 preds는 tensor의 가장 큰 index 값, 즉 확률이 가장 큰 class 값이므로
tensor([0, 1, 2, 2], device='cuda:0') <- 이런 값을 갖는다.

최종 결과