Pytorch로 이미지 분류하기

Pytorch를 이용해 이미지 데이터를 분류 해보도록 하겠습니다.

실습을 위해 사용한 데이터는 데이콘 에서 다운 받으실 수 있습니다.

해당 대회는 28 x 28 pixel 크기의 손글씨 이미지가 어떤 숫자를 담고 있는지 분류하는 것이 목표입니다.

 

라이브러리 불러오기

우선 실습에 필요한 라이브러리를 불러오도록 하겠습니다.

import warnings
warnings.filterwarnings(action='ignore')

import os
import pandas as pd
import numpy as np

from PIL import Image
import matplotlib.pyplot as plt
%matplotlib inline

from sklearn.model_selection import train_test_split
import IPython.display as ipd
import time

import torch
import torchvision
from torch import Tensor, nn, optim
from torch.utils.data import Dataset, DataLoader

import torch.nn.functional as F
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.models as models
from adamp import AdamP
from torchsummary import summary

import albumentations as A
from albumentations.pytorch import ToTensor
from efficientnet_pytorch import EfficientNet

 

시드 고정

시드를 고정해 모델의 추론 값이 일정하게 유지 될 수 있도록 해줍니다. 

시드를 고정하지 않으면 커널을 다시 시작하여 모델을 학습 할 때 마다 결과값이 다를 수 있습니다. 

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

 

데이터 불러오기

판다스의 read_csv를 이용해 데이터를 불러옵니다.

해당 대회의 train.csv는 이미지 파일의 이름과 분류 값인 label로 이뤄져 있습니다. 

또한 test.csv는 이미지 파일의 이름값으로만 이뤄져 있습니다.

train = pd.read_csv('./train/train_data.csv')
test = pd.read_csv('./test/test/test_data.csv')

train.columns = ['file_name','label'] # 컬럼명 변경

train.head()

 

커스텀 데이터셋 정의

PyTorch의 mini batch 학습에 필요한 커스텀 데이터셋을 정의 해줍니다. 

class CustomDataset(torch.utils.data.Dataset): 
    def __init__(self,df,path,option,augmentation=None):
        self.df = df
        self.option = option
        self.augmentation = augmentation
        self.path = path

    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx): 
        file_path  = self.df.iloc[idx,0]
        image=  Image.open(os.path.join(self.path,file_path)).convert('RGB')
        
        image = np.array(image)
             
        
        
        if self.augmentation is not None:
            image = self.augmentation(image=image)['image']
        
        
        if self.option =='train':
            label = self.df.iloc[idx,1]
            label = torch.tensor(label, dtype=torch.int64)
            return image, label
        
        return image

 

모델 정의

모델을 정의 해줍니다. 이번 실습에서는 사전학습 되지 않은 resnet 구조를 사용합니다.

class MnistModel(nn.Module):
    def __init__(self,):
        super().__init__()
        self.model = torchvision.models.resnet18(pretrained = False)
        #self.model = EfficientNet.from_name('efficientnet-b1') efficient net 사용
        self.dropout = nn.Dropout(0.5)
        
        self.classifier = nn.Linear(1000, 10)

    def forward(self, images):
        outputs = self.model(images)
        outputs = self.dropout(outputs)
        outputs = self.classifier(outputs)
        return outputs

 

모델 학습

모델 학습 위한 파라미터를 정의하고 학습을 진행합니다. 

또한 k-fold를 이용하여 교차 검증도 진행합니다.

# k-fold 정의
from sklearn.model_selection import StratifiedKFold

skf = StratifiedKFold(n_splits = 4)
folds=[]
for trn_idx,val_idx in skf.split(train['file_name'],train['label']):
    folds.append((trn_idx,val_idx))

num_epochs = 30 # 학습 횟수
best_models = [] # 폴드별로 가장 validation acc가 높은 모델 저장

for i,fold in enumerate(range(4)):
    print('===============',i+1,'fold start===============')
    
    model = CNN().to('cuda') 
    optimizer = AdamP(model.parameters(), lr=1e-3 ) # optimizer 정의
    criterion = nn.CrossEntropyLoss() # 손실함수 정의 
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size = 5,
                                                gamma = 0.9) # learning rate scheduler 정의
        
    
    train_idx = folds[fold][0] #
    valid_idx = folds[fold][1]
    
    train_data = train.loc[trn_idx]
    val_data = train.loc[valid_idx]
    
    # 데이터 증강
    train_transforms = A.Compose([
    A.Normalize(),
    A.HorizontalFlip(),
    A.RandomRotate90(p=0.5),
    A.VerticalFlip(p=0.5),
    ToTensor()
    ])

    test_transforms = A.Compose([
        A.Normalize(),
        ToTensor()
        ])

    train_dataset = CustomDataset(train_data,'train/','train',train_transforms)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)

    valid_dataset = CustomDataset(val_data,'train/','train',test_transforms)
    valid_loader = DataLoader(valid_dataset, batch_size=64, shuffle=False, num_workers=0)
    
    
    valid_acc_max = 0.85
    valid_loss_min = 0.4
    
    for epoch in range(num_epochs):
		
        #모델 학습
        for i, (images, targets) in enumerate(train_loader):
            model.train()
            images, targets = images.to('cuda'), targets.to('cuda')

            optimizer.zero_grad()
            outputs = model(images)

            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()  

            if (i+1) % 20 == 0:
                print(f'Epoch: {epoch} - Loss: {loss:.6f}')


        val_loss = []
        val_acc = []
        # 모델 검증
        for i, (images, targets) in enumerate(valid_loader):
            model.eval()
            images, targets = images.to('cuda'), targets.to('cuda')

            with torch.no_grad():
                outputs = model(images)
                valid_loss = criterion(outputs,targets).cpu().detach().numpy()

                preds = torch.argmax(outputs,axis = 1)
                preds = preds.cpu().detach().numpy()

                targets = targets.cpu().detach().numpy()
                batch_acc = (preds==targets).mean()



                val_loss.append(valid_loss)
                val_acc.append(batch_acc)

        val_loss = np.mean(val_loss)
        val_acc = np.mean(val_acc)

        print(f'Epoch: {epoch} - valid Loss: {val_loss:.6f} - valid_acc : {val_acc:.6f}')
        
        # 성능 좋은 모델을 저장하는 부분 이후 예측시 사용 
        # 위 코드는 정확도 기준 아래 코드는 loss 기준
        '''
         if valid_acc_max < val_acc:
            valid_acc_max = val_acc
            best_models.append(model)
            print('model save, model val acc : ',val_acc)
            print('best_models size : ',len(best_models))
            '''

		
        if valid_loss_min > val_loss:
            valid_loss_min = val_loss
            best_models.append(model)
            
    # Learning rate 조절
    lr_scheduler.step()

 

test 데이터 예측

학습한 모델을 이용하여 test 데이터를 예측 합니다.

test_dataset = CustomDataset(test,'test/test','test',test_transforms)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)

preds = []
for idx,model in enumerate(best_models[:1]): 
    print(idx+1, '번째 모델 예측 진행중')
    model = model
    model.eval()
    y_pred = []
    with torch.no_grad():
        for i, d in enumerate(test_loader):
            inputs = d.to('cuda')
            outputs = model(inputs).detach().cpu().numpy()
            y_pred.extend(outputs.argmax(axis=1).astype(int))
            
    preds.append(y_pred)

 

 

예측결과 hard voting

from collections import Counter
np_pred = np.array(preds).T

pred = []
for i in range(5000):
    cnt = Counter(np_pred[i])
    pred.append(cnt.most_common()[0][0])
#pred

 

이상으로 pytorch를 이용한 이미지 분류 코드에 대한 실습을 마무리 하겠습니다.

궁금하신 점은 댓글로 남겨주세요 !