[pytorch] Image Classification : 🐶 (Using ViT)

지난 시간에 이어, Dataset / DataLoader class를 활용하여 강아지 분류 모델을 작성해보도록 하겠습니다! 😎

✨ 왜 강아지냐구요? 🐶 강아지는 귀여우니까요 :)

😉 그럼, LET'S DIGGIN' !

Data 준비

✔ Data는 Kaggle에서 Stanford Dog Dataset을 준비했습니다.
✔ https://www.kaggle.com/datasets/jessicali9530/stanford-dogs-dataset

• 800MB 정도 나가고, 라벨도 120개 정도로 적당해 보이더라구요 :)

✔ train : val = 0.85 : 0.15 split 수행하였습니다.

import os
import shutil

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
org_image_path = os.path.join(root_path, "archive/images/Images/")
labels = os.listdir(org_image_path)

labels_cnt_list = []

for l in labels:
    labels_cnt_list.append(len(os.listdir(os.path.join(org_image_path, l))))

os.makedirs(os.path.join(root_path, 'data/'), exist_ok=True)
os.makedirs(os.path.join(root_path, 'data/train/'), exist_ok=True)
os.makedirs(os.path.join(root_path, 'data/val/'), exist_ok=True)

for l in labels:
    os.makedirs(os.path.join(root_path, 'data/train/', l), exist_ok=True)
    os.makedirs(os.path.join(root_path, 'data/val/', l), exist_ok=True)

train_img_list = []
train_label_list = []
val_img_list = []
val_label_list = []

for idx, l in enumerate(labels):
    num_train = int(labels_cnt_list[idx] * 0.85)
    tmp_image_name_list = os.listdir(os.path.join(org_image_path, l))
    for cnt, fname in enumerate(tmp_image_name_list):
        if cnt <= num_train:
        	dst_path = os.path.join(root_path, 'data/train/', l, fname)
            shutil.copy(os.path.join(org_image_path, l, fname), dst_path)
            train_img_list.append(dst_path)
            train_label_list.append(idx)
        else:
        	dst_path = os.path.join(root_path, 'data/val/', l, fname)
            shutil.copy(os.path.join(org_image_path, l, fname), dst_path)
            val_img_list.append(dst_path)
            val_label_list.append(idx)

ViT?

✔ Vit : Vision Transformer!
✔ Computer Vision Task에서 항상 빠지지 않던 CNN 아키텍처를 제외하고, 오직 Self-attention만을 사용하여서도 충분히 CV Task를 수행할 수 있음을 보여준 논문입니다 :)
✔ pytorch에 구현체가 준비되어 있어, 그대로 활용하면 됩니다.
✨ timm (PyTorch Image Models) 패키지를 사용할 예정입니다!
😎 timm 패키지는 높은 성능을 보이는 Computer Vision D/L 알고리즘들을 미리 pytorch로 구현해 놓은 패키지로서, 간단하고 빠르게 모델을 구현할 수 있게 도와줍니다 :)

pip install timm

https://arxiv.org/abs/2010.11929

Dataset code

지난 시간에 구현하였던 Dataset과 DataLoader class를 활용하겠습니다. 😉

Dataset class 선언

import torch
from PIL import Image

class MyDataset(torch.utils.data.Dataset):
    """
    Attributes
    ----------
    img_list : 리스트
        이미지의 경로를 저장한 리스트
    label_list : 리스트
        label의 경로를 저장한 리스트
    phase : 'train' or 'val'
        학습 또는 테스트 여부 결정
    transform : object
        전처리 클래스의 인스턴스
    """

    def __init__(self, img_list, label_list, phase, transform):
        self.img_list = img_list
        self.label_list = label_list
        self.phase = phase  # train 또는 val을 지정
        self.transform = transform  # 이미지의 변형

    def __len__(self):
        '''이미지의 갯수를 반환'''
        return len(self.img_list)

    def __getitem__(self, index):
        '''
        전처리한 이미지 및 라벨 return
        '''
        # img_path = self.img_list[index]
        # img = Image.open(img_path).convert('RGB')
        img = self.img_list[index]
        
        transformed_img = self.transform(img, self.phase)
        label = self.label_list[index]
        
        return transformed_img, label

Transform class 선언

from torchvision import models, transforms

class MyTransform():
    """
    Attributes
    ----------
    resize : int
        Transform 수행 후 변경될 width / height 값.
    mean : (R, G, B)
        각 색상 채널의 평균값.
    std : (R, G, B)
        각 색상 채널의 표준 편차.
    """

    def __init__(self, resize, mean, std):
        self.data_transform = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(
                    (resize, resize), scale=(0.5, 1.0)),  
                transforms.RandomHorizontalFlip(), 
                transforms.ToTensor(),  # 텐서로 변환
                transforms.Normalize(mean, std)  # 표준화
            ]),
            'val': transforms.Compose([
                transforms.Resize((resize, resize)),
                transforms.ToTensor(),  # 텐서로 변환
                transforms.Normalize(mean, std)  # 표준화
            ])
        }

    def __call__(self, img, phase='train'):
        """
        Parameters
        ----------
        phase : 'train' or 'val'
            전처리 모드를 지정.
        """
        return self.data_transform[phase](img)

😎 좋습니다! 이제 사전 선언해야 할 클래스는 모두 선언하였습니다.
😏 그러면, 학습을 수행하는 코드를 작성해보도록 하겠습니다.

Train Code

Dataset instance 생성

size = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)

train_dataset = MyDataset(img_list=train_img_list, label_list=train_label_list, phase="train", transform=MyTransform(
    size, mean, std)))

val_dataset = MyDataset(img_list=val_img_list, label_list=val_label_list, phase="val", transform=MyTransform(
    size, mean, std)))

image_datasets = {'train' : train_dataset, 'val' : val_dataset}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}

DataLoader instance 생성

batch_size = 32

train_dataloader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)

val_dataloader = torch.utils.data.DataLoader(
    val_dataset, batch_size=batch_size, shuffle=False)

# 사전 객체에 정리
dataloaders_dict = {"train": train_dataloader, "val": val_dataloader}

ViT model

import timm

num_classes = 120
model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=num_classes)

loss function & Optimizer

import torch.optims as optims
import torch.nn as nn
from torch.optim import lr_scheduler

citerion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

자, 진짜 학습 시작입니다! :)

Train loop 작성

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

epochs = 10

for epoch in range(epochs):
	print("{}/{} epoch running now".format(epoch, epochs - 1))
    
    for phase in ['train', 'val']:
    	if phase == 'train':
        	model.train()
        else:
        	model.eval()
        
        running_loss = 0.0
        running_corrects = 0
        
        for inputs, labels in dataloaders[phase]:
        	inputs = inputs.to(device)
            labels = labels.to(device)
            
            optimizer.zero_grad()
            
            # requires_grad = True, when Training
            with torch.set_grad_enabled(phase == 'train'):
            	outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)
                
                if phase == 'train':
                	# 손실 계산 (각 Tensor의 gradient 계산)
                	loss.backward()
                    # 계산된 손실을 optimizer의 계산에 따라 weight 조정
                    optimizer.step()
                    
            running_loss += loss.item() * inputs.size(0)
            running_corrects = torch.sum(preds == labels.data)
        
        if phase == 'train':
        	scheduler.step()
            
        epoch_loss = running_loss / dataset_sizes[phase]
        epoch_acc = running_corrects.double() / dataset_sizes[phase]
        
        print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
		
        # 모델 저장
        if phase == 'val' and epoch_acc > best_acc :
        	best_acc = epoch_acc
            torch.save('./best_model.pth')

😎 이로써 강아지 데이터를 활용하여 ViT 모델의 Transfer Learning을 간단히 실습하는 코드를 작성해 보았습니다.
😏 하지만 논문에 명시된 Optimizer나, learning rate scheduler는 다른 부분이 있기에, 이 부분은 추후 개선 사항이 되겠군요 :)

Visualization Code

https://tutorials.pytorch.kr/beginner/transfer_learning_tutorial.html

👍 시각화 코드는 위의 튜토리얼 페이지의 함수를 사용하였습니다.

def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure()

    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)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis('off')
                ax.set_title('predicted: {}'.format(class_names[preds[j]]))
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)

😉 아래와 같이 사용하면 됩니다.

visualize_model(model)

마치며

🎉 WOW! 감동적이게도 뛰어난 성능을 보여준 모델을 가져와서 Transfer Learning을 수행하는 데 성공하였습니다!
😋 물론, 논문의 벤치마크를 재현하려면 논문 그대로의 Training sceinaro와 Optimizer, 그리고 learning rate scheduling 등이 포함되어야 하지만, 일단 해낸게 어디에요 :)
😁 다음 기회에는 좀더 Advanced 한 구현으로 도전해보도록 하겠습니다!