[Pytorch][Kaggle] Cats vs. Dogs Classification

(torch version : 1.7.0)

지난번 tensorflow에서 Cats vs. Dogs Classification에 이어서, Pytorch를 사용해서 해당 문제를 구현해보도록 하겠습니다.

2020/12/04 - [ML & DL/tensorflow] - [Tensorflow][Kaggle] Cats vs. Dogs Classification

[Tensorflow][Kaggle] Cats vs. Dogs Classification

tensorflow에서는 VGG16 모델을 fine tuning해서 진행했었는데, pytorch에서는 torchvision에서 제공하는 ResNet50을 사용해서 진행합니다

 

 사용되는 package는 다음과 같습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import PIL
import shutil
import zipfile
import glob
import os
import time

 

1. Dataset 준비

tensorflow의 경우에는 ImageDataGenerator를 사용하기 위해 각 class를 폴더별로 분류했지만, 이번에는 training data를 train/valid/test set별로 폴더를 나누어서 이미지를 저장합니다.

data_zip_dir = '/kaggle/input/dogs-vs-cats-redux-kernels-edition'
train_zip_dir = os.path.join(data_zip_dir, 'train.zip')
test_zip_dir = os.path.join(data_zip_dir, 'test.zip')

with zipfile.ZipFile(train_zip_dir, 'r') as z:
    z.extractall()
with zipfile.ZipFile(test_zip_dir, 'r') as z:
    z.extractall()

우선 input data의 경로를 설정하고, 압축을 해제합니다.

 

train_dir = os.path.join(os.getcwd(), 'train')
test_dir = os.path.join(os.getcwd(), 'test')

train_set_dir = os.path.join(train_dir, 'train')
os.mkdir(train_set_dir)
valid_set_dir = os.path.join(train_dir, 'valid')
os.mkdir(valid_set_dir)
test_set_dir = os.path.join(train_dir, 'test')
os.mkdir(test_set_dir)

dog_files = [f'dog.{i}.jpg' for i in range(12500)]
cat_files = [f'cat.{i}.jpg' for i in range(12500)]

그리고 각 용도에 맞는 폴더를 생성합니다.

이미지 파일의 이름이 cat.{number}.jpg / dog.{number}.jpg의 형태로 되어있는데, 각 파일 명을 list로 저장해둡니다.

 

for dog, cat in zip(dog_files[:10000], cat_files[:10000]):
    src = os.path.join(train_dir, dog)
    dst = os.path.join(train_set_dir, dog)
    shutil.move(src, dst)
    src = os.path.join(train_dir, cat)
    dst = os.path.join(train_set_dir, cat)
    shutil.move(src, dst)
    
for dog, cat in zip(dog_files[10000:11250], cat_files[10000:11250]):
    src = os.path.join(train_dir, dog)
    dst = os.path.join(valid_set_dir, dog)
    shutil.move(src, dst)
    src = os.path.join(train_dir, cat)
    dst = os.path.join(valid_set_dir, cat)
    shutil.move(src, dst)
    
for dog, cat in zip(dog_files[11250:12500], cat_files[11250:12500]):
    src = os.path.join(train_dir, dog)
    dst = os.path.join(test_set_dir, dog)
    shutil.move(src, dst)
    src = os.path.join(train_dir, cat)
    dst = os.path.join(test_set_dir, cat)
    shutil.move(src, dst)

하나의 폴더에 존재하는 이미지 파일들을 적절하게 분배해서 train/valid/test 폴더에 각각 옮겨줍니다.

print(f'the number of train set : {len(os.listdir(train_set_dir))}')
print(f'the number of validn set : {len(os.listdir(valid_set_dir))}')
print(f'the number of test set : {len(os.listdir(test_set_dir))}')

train/valid/test용으로 각 20,000/2,500/2,500개의 이미지로 분배했습니다.

 

2. Data Generator 생성

tensorflow와 유사한 방법으로 DataLoader를 생성해서 학습에 사용해보도록 하겠습니다.

현재 데이터에 맞게 DataLoader를 사용하기 위해서는 우선 CustomDataset 클래스를 생성해야하는데, torch.utils.data.Dataset 클래스를 상속받아서 용도에 맞는 Datatset 클래스를 생성합니다. 

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, files, root, mode='train', transform=None):
        self.files = files
        self.root = root
        self.mode = mode
        self.transform=transform
        
        if 'cat' in files[0]:
            self.label = 0
        else:
            self.label = 1
    
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, index):
        img = PIL.Image.open(os.path.join(self.root, self.files[index]))
        
        if self.transform:
            img = self.transform(img)
        if self.mode == 'train':
            return img, np.array([self.label])
        else:
            return img, self.files[index]

torch.utils.data.Dataset을 상속받는 클래스는 위와 같이 구현할 수 있습니다. 그리고 이렇게 상속받은 클래스는 __getitem__과 __len__를 정의해서 내부 동작을 구현할 수 있습니다.

 

클래스 함수를 간단하게 살펴보면 다음과 같습니다.

- __init__(self, files, root, mode='train', transform=None) :

• files : 이미지 파일 이름을 저장하고 있는 list
• root : 이미지 파일이 존재하는 폴더 경로
• mode : 해당 dataset이 train용인지 eval용인지 체크
• transform : 이미지의 전처리를 위한 torchvision.transform

- __len__(self) : Dataset의 길이를 반환하기 위한 메소드

- __getitem(self, index) : 주어진 key에 해당하는 data를 반환하는 메소드이며, key에 해당하는 이미지 파일을 읽고, 전처리 과정을 통해서 data를 리턴합니다. mode='train'일 경우에는 label을 반환하고, 'train'용이 아닌 경우에는 label을 모르기 때문에 실제 이미지 파일의 경로를 반환하도록 하였습니다.

 

다음은 이미지 전처리에 사용되는 torchvision.transform 입니다.

train_transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((256,256)),
    torchvision.transforms.RandomCrop(224),
    torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor(),
])
test_transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224,244)),
    torchvision.transforms.ToTensor(),
])

train용으로는 우선 (256,256)의 사이즈로 변환하고, RandomCrop을 적용해서 Zoom 효과를 적용하였고, RandomHorizontalFlip을 통해서 랜덤하게 좌우 반전 효과도 주었습니다.

(tensorflow처럼 기울기나 여러가지 효과를 적용할 수도 있습니다.) 

평가를 위한 transform은 다른 효과를 주면 안되기 때문에 이미지의 사이즈만 동일하게 맞추어주었습니다. 

마지막으로는 ToTensor()를 적용하였는데, 이것은 기본적으로 PIL image나 numpy array를 torch.FloatTensor로 변환하고 [0,255] 범위를 갖는 픽셀값을 [0,1]의 범위로 변환까지 해줍니다.

scaling이 항상 적용되는 것은 아니며, PIL image가 (L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK, 1) mode중의 하나이거나, numpy array의 dtype이 np.unit8인 경우에 scaling이 적용됩니다.

 

train_dog_dataset = CustomDataset(dog_files[:10000], train_set_dir, transform=train_transform)
train_cat_dataset = CustomDataset(cat_files[:10000], train_set_dir, transform=train_transform)
valid_dog_dataset = CustomDataset(dog_files[10000:11250], valid_set_dir, transform=test_transform)
valid_cat_dataset = CustomDataset(cat_files[10000:11250], valid_set_dir, transform=test_transform)
test_dog_dataset = CustomDataset(dog_files[11250:], test_set_dir, transform=test_transform)
test_cat_dataset = CustomDataset(cat_files[11250:], test_set_dir, transform=test_transform)

train_dataset = torch.utils.data.ConcatDataset([train_dog_dataset, train_cat_dataset])
valid_dataset = torch.utils.data.ConcatDataset([valid_dog_dataset, valid_cat_dataset])
test_dataset = torch.utils.data.ConcatDataset([test_dog_dataset, test_cat_dataset])

CustomDataset을 정의했다면, 용도별/라벨별로 dataset을 생성해주고 용도별로 합쳐주는 작업까지 진행해줍니다.

print(f'number of train dataset : {len(train_dataset)}')
print(f'number of valid dataset : {len(valid_dataset)}')
print(f'number of test dataset : {len(test_dataset)}')

마지막으로 학습과 평가에 사용할 DataLoader들을 생성해줍니다.

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=True)

 

DataLoader가 제대로 적용이 되는지 확인을 해보도록 합시다.

samples, labels = iter(train_loader).next()
classes = {0:'cat', 1:'dog'}
fig = plt.figure(figsize=(16,24))
for i in range(24):
    a = fig.add_subplot(4,6,i+1)
    a.set_title(classes[labels[i].item()])
    a.axis('off')
    a.imshow(np.transpose(samples[i].numpy(), (1,2,0)))
plt.subplots_adjust(bottom=0.2, top=0.6, hspace=0)

정상적으로 이미지들이 로드되는 것을 확인할 수 있습니다.

 

3. ResNet Model Load 및 Prediction layer 수정

GPU를 사용해서 학습을 진행할 예정이므로, device 변수를 우선 설정해줍니다. (cuda로 설정되어야 합니다)

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

model = torchvision.models.resnet50(pretrained=True)
model

torchvision.models에서 제공하는 pretrained된 resnet50을 사용하도록 하겠습니다. 그리고 우리는 개와 고양이 분류를 위해 binary classification을 진행해야하지만, 모델의 마지막 부분을 살펴보면 1000개의 feature를 출력으로 보내고 있습니다.

따라서 이 부분을 수정해보도록 하겠습니다. 수정 방법은 간단합니다.

num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
    nn.Dropout(0.5),
    nn.Linear(num_ftrs, 1024),
    nn.Dropout(0.2),
    nn.Linear(1024, 512),
    nn.Dropout(0.1),
    nn.Linear(512, 1),
    nn.Sigmoid()
)

nn.Sequential을 마지막 출력이 1이 되도록(Sigmoid 연산까지 추가해줍니다) 구성하고, fc에 할당해줍니다.

그런 다음, 다시 확인해보면 변경된 것을 볼 수 있습니다.

 

추가적으로 tensorflow의 medel.summary()와 같이 torchsummary를 통해서 모델의 요약정보를 확인할 수 있습니다. kaggle에 torchsummary가 설치되어 있지 않기 때문에, 우선 설치를 해주고 import 합니다.

!pip install torchsummary
from torchsummary import summary

그리고 모델의 요약정보를 다음과 같은 방법을 통해서 확인할 수 있는데, torchsummary는 모델의 GPU에 로드되어 있어야 사용이 가능해서, cuda()를 통해서 GPU에 로드시킨 후에 확인할 수 있습니다.(input_size도 입력해주어야 합니다.)

model.cuda()
summary(model, input_size=(3,224,224))

 

이제 학습에 사용할 모델 구성을 완료하였습니다.

 

4. Training

학습을 진행하기 전, 학습에 필요한 fit 함수를 정의하고 fit 함수를 구현해서 학습을 진행해보도록 하겠습니다.

def fit(model, criterion, optimizer, epochs, train_loader, valid_loader):
    model.train()
    
    train_loss = 0
    train_acc = 0
    train_correct = 0
    
    train_losses = []
    train_accuracies = []
    valid_losses = []
    valid_accuracies = []
    
    for epoch in range(epochs):
        start = time.time()
        for train_x, train_y in train_loader:
            model.train()
            train_x, train_y = train_x.to(device), train_y.to(device).float()
            optimizer.zero_grad()
            pred = model(train_x)
            loss = criterion(pred, train_y)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            
            y_pred = pred.cpu()
            y_pred[y_pred >= 0.5] = 1
            y_pred[y_pred < 0.5] = 0
            train_correct += y_pred.eq(train_y.cpu()).int().sum()
        
        # validation data check
        valid_loss = 0
        valid_acc = 0
        valid_correct = 0
        for valid_x, valid_y in valid_loader:
            with torch.no_grad():
                model.eval()
                valid_x, valid_y = valid_x.to(device), valid_y.to(device).float()
                pred = model(valid_x)
                loss = criterion(pred, valid_y)
            valid_loss += loss.item()
            
            y_pred = pred.cpu()
            y_pred[y_pred >= 0.5] = 1
            y_pred[y_pred < 0.5] = 0
            valid_correct += y_pred.eq(valid_y.cpu()).int().sum()
        
        train_acc = train_correct/len(train_loader.dataset)
        valid_acc = valid_correct/len(valid_loader.dataset)
        
        print(f'{time.time() - start:.3f}sec : [Epoch {epoch+1}/{epochs}] -> train loss: {train_loss/len(train_loader):.4f}, train acc: {train_acc*100:.3f}% / valid loss: {valid_loss/len(valid_loader):.4f}, valid acc: {valid_acc*100:.3f}%')
        
        train_losses.append(train_loss/len(train_loader))
        train_accuracies.append(train_acc)
        valid_losses.append(valid_loss/len(valid_loader))
        valid_accuracies.append(valid_acc)
        
        train_loss = 0
        train_acc = 0
        train_correct = 0
    
    plt.plot(train_losses, label='loss')
    plt.plot(train_accuracies, label='accuracy')
    plt.legend()
    plt.title('train loss and accuracy')
    plt.show()
    
    plt.plot(valid_losses, label='loss')
    plt.plot(valid_accuracies, label='accuracy')
    plt.legend()
    plt.title('valid loss and accuracy')
    plt.show()

필요한 정보들을 다 담아보려고 했는데, 꽤나 길어진것 같군요.

fit함수는 model, loss function, optimizer와 epoch, train_loader, valid_loader를 인자로 받고 있습니다.

설정한 epoch만큼 학습을 진행하고, 매 epoch가 끝나면 validation check를 진행하도록 했습니다. 

설정한 epoch만큼 학습이 끝난 후에는 loss와 accuracy의 변화 추이를 살펴보기 위해서 매 epoch마다 loss, acc를 저장하고 있으며, acc는 한 epoch에서 label과 일치하는 갯수를 저장해서 총 data 갯수로 나누어서 계산하게 됩니다.

 

모델의 output은 sigmoid를 거쳐서, label class가 1인 확률이 됩니다.(input image가 dog일 확률) 따라서, 0.5이상은 1, 0.5미만은 0으로 처리해주고, 실제 label과 비교해서 일치하는지 확인합니다.

 

그럼 이제 loss function과 optimizer를 생성하고 학습을 진행해보도록 하겠습니다.

Loss function으로는 Binary Cross Entropy Loss(BCELoss())를 사용했고, Optimizer는 Adam을 사용했으며, ResNet을 fine turning하기 때문에 learning rate를 0.00001로 낮게 설정했습니다.

criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

fit(model, criterion, optimizer, 10, train_loader, valid_loader)

10 epoch도 조금 오버였던 것 같습니다.. 1~2 epoch만으로도 충분히 학습이 된 것으로 보이네요.

 

5. 모델 평가

이전에 만든 test_loader를 통해서, 학습한 모델을 평가해보도록 하겠습니다.

평가에 사용할 eval 함수입니다.

def eval(model, criterion, test_loader):
    with torch.no_grad():
        model.eval()
        correct = 0
        losses = 0
        for test_x, test_y in test_loader:
            test_x, test_y = test_x.to(device), test_y.to(device).float()
            pred = model(test_x)
            loss = criterion(pred, test_y)
            
            y_pred = pred.cpu()
            y_pred[y_pred >= 0.5] = 1
            y_pred[y_pred < 0.5] = 0
            
            losses += loss.item()
            correct += y_pred.eq(test_y.cpu()).int().sum()
    print(f'eval loss: {losses/len(test_loader):.4f}, eval acc: {correct/len(test_loader.dataset)*100:.3f}%')

eval(model, criterion, test_loader)

99%의 정확도를 얻고 있습니다.

 

6. 결과 제출

이렇게 학습한 모델을 가지고, 이제 test.zip의 데이터를 사용해서 결과를 제출해보도록 하겠습니다.

submit_files = [f'{i}.jpg' for i in range(1, 12500+1)]
submit_dataset = CustomDataset(submit_files, test_dir, mode='test', transform=test_transform)
submit_loader = torch.utils.data.DataLoader(submit_dataset, batch_size=32, shuffle=False)

제출용 파일은 {숫자}.jpg의 형태를 띄고 있습니다. 총 12,500개의 이미지가 있으므로 1에서 12500까지의 파일 이름 list를 생성해주고, DataLoader를 만들어줍니다. os.listdir(test_dir)를 통해서 파일 이름을 읽어올 수도 있지만, 이렇게 읽어올 경우에 파일 이름이 정렬이 되어 있지 않습니다.

 

제대로 생성이 되었는지 확인해보도록 하겠습니다.

imgs, files = iter(submit_loader).next()
fig = plt.figure(figsize=(16,24))
for i in range(24):
    a = fig.add_subplot(4,6,i+1)
    a.set_title(files[i])
    a.axis('off')
    a.imshow(np.transpose(samples[i].numpy(), (1,2,0)))
plt.subplots_adjust(bottom=0.2, top=0.6, hspace=0)

잘 읽어오는 것 같네요

 

이제 predict 함수를 통해서 확률을 나타내는 결과를 받아오도록 하겠습니다.

def predict(model, data_loader):
    with torch.no_grad():
        model.eval()
        ret = None
        for img, files in data_loader:
            img = img.to(device)
            pred = model(img)
            
            if ret is None:
                ret = pred.cpu().numpy()
            else:
                ret = np.vstack([ret, pred.cpu().numpy()])
    return ret
   
pred = predict(model, submit_loader)

 

샘플을 뽑아서 제대로 예측하고 있는지 확인해보겠습니다.

sample_pred = pred[:24]
sample_pred[sample_pred >= 0.5] = 1
sample_pred[sample_pred < 0.5] = 0

imgs, files = iter(submit_loader).next()
classes = {0:'cat', 1:'dog'}
fig = plt.figure(figsize=(16,24))
for i in range(24):
    a = fig.add_subplot(4,6,i+1)
    a.set_title(classes[sample_pred[i][0]])
    a.axis('off')
    a.imshow(np.transpose(imgs[i].numpy(), (1,2,0)))
plt.subplots_adjust(bottom=0.2, top=0.6, hspace=0)

샘플 데이터는 모두 다 맞춘 것 같네요 !

 

이제 csv 파일로 변환해서 제출해보도록 하겠습니다.

submission = pd.DataFrame(np.clip(pred, 1e-6, 1-1e-6), columns=['label'])
submission['id'] = submission.index + 1
submission = submission[['id', 'label']]
submission.head(10)
submission.to_csv('submission.csv', index=False)

확률이 너무 작은값이 되지 않도록 cliping 처리를 해주고, csv 파일로 생성하고 제출해주면 끝입니다.

0.07의 준수한 결과를 얻었습니다 ! (leaderboard의 1위가 0.033입니다)

확실히 pre-trained된 모델을 사용하게 되면 시간은 줄이고, 성능은 대폭 향상시킬 수 있는 것 같습니다.