ImageNet Pretrained ResNet 으로 MNIST 분류기 만들기

1. 전이 학습(Transfer Learning)이란?

위키백과의 정의에 따르면, '한 분야의 문제를 해결하기 위해 얻은 지식과 정보를 단른 문제를 푸는데 사용하는 방식"이다. 딥러닝 분야에서는, 이미지분류 문제를 해결하는데 사용헀떤 네트워크를 다른 데이터셋 혹은 다른 문제(Task)에 적용시켜 푸는 것을 의미한다. 특히, 시각적 이해를 목표로 하는 컴퓨터 비전의 영역에서 전이학습으로 수행된 모델들이 높은 성능을 보이고 있어, 가장 많이 사용되는 대표적인 방법중 하나이다. 

 

이러한 결과를 보여주는 이유는, 다양한 이미지의 "특징"(feature)들을 학습했기 때문이다. 일반적으로 네트워크가 깊어질 수록, 서로 다른 종류들의 피처들을 학습한다고 알려져있는데, 낮은 층에서 학습되는 feature를 low-level-feature, 깊은 층에서 학습되는 feature들은, high-level-feature라고 부른다. low-level-feature의 예로는, 이미지의 색이나 경계(edge)등을 말할 수 있고, high-level-feature는 더 심화된 객체의 패턴이나 형태를 의미한다. 

(low-level-feature 는 주로 local 한 feature(색의 변화나 경계의 방향등)를, high-level-feature 는 global 한 feature(동그라미가 반복되는 패턴 또는 새의 부리등)를 학습한다고 알려져있다 ) 

 

2. 전이 학습 사용 방법

만약 ImageNet과 비슷하지만 소량의 데이터셋을 가지고 있다면 ImageNet으로 학습시킨 CNN을 구조를 그대로 두고 뒷단에 분류를 위해 새로운 완전연결레이어(FC;Fully Connected Layer)를 붙여서 학습시키면 된다. 만약 ImageNet과 비슷하지만 더 많은 데이터셋을 가지고 있다면 뒷단에 여러 FC레이어를 묶어서 학습시켜도 좋다고 한다.

 

만약 내가 가진 데이터셋이 ImageNet과 다른 경우에는 어떨까? 데이터가 많은 경우에 뒷단의 여러 FC 레이어를 묶어서 학습시키는 것으로 충분하다. 하지만 데이터가 적은 경우엔 pre-training이 효과적이지 않을 수 있으며, 이러한 경우, data augmentation으로 데이터의 양을 늘려서 전체 네트워크를 학습시키는 등 다른 방법을 생각해야한다.

 

3. 전이학습의 효과

2018년 FAIR(Facebook AI Research) 논문에서, 실험을 통해 '전이학습이 학습속도면에서 효과가 있음'을 보여주었다. 동일한 조건에서, 기본 모델을 사용하는 것보다, 사전학습된 모델을 사용하는 것이, 짧게 시간이 걸림을 확인할 수 있다. ( y축 accuarcy x 축 시간 )

 

 

https://dacon.io/forum/405988

전이학습(Transfer Learning)이란?

해당 내용을 참조하였습니다. 

 

---

4. ImageNet Pretrained ResNet 으로 MNIST 분류기 만들기

2023.08.15 - [인공지능공부] - CNN을 활용한 MNIST 분류 모델 구현

CNN을 활용한 MNIST 분류 모델 구현

코드에 대한 자세한 설명은 위 포스팅 참고.

라이브러리 불러오기 부터, confusion matrix 만드는 코드 , 성능 평가하는 코드는 동일함. (모델 부분만 다름)

 

1. 라이브러리 불러오기

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import torchvision
import torchvision.transforms as transforms

from matplotlib import pyplot as plt
import seaborn as sn
import pandas as pd
import numpy as np
import os

device = 'cuda'

 

2. 데이터 준비

- Pretrained 된 모델의 경우, 입력 데이터 크기가 224 x 224 로 고정되어 있으므로, transforms.Resize로 크기 조절해야함.

# 사전 학습된 모델의 경우 입력 데이터의 크기가 224 x 224이므로, MNIST 데이터셋의 크기를 조절할 필요가 있음
transfer_transform_train = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
])

transfer_transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
])

transfer_train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transfer_transform_train)
transfer_test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transfer_transform_test)

transfer_train_loader = torch.utils.data.DataLoader(transfer_train_dataset, batch_size=128, shuffle=True, num_workers=4)
transfer_test_loader = torch.utils.data.DataLoader(transfer_test_dataset, batch_size=100, shuffle=False, num_workers=4)

 

3. 학습(Training) 및 평가(Testing) 함수 정의

def train(net, epoch, optimizer, criterion, train_loader):
    print('[ Train epoch: %d ]' % epoch)
    net.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()

        benign_outputs = net(inputs)
        loss = criterion(benign_outputs, targets)
        loss.backward()

        optimizer.step()
        train_loss += loss.item()
        _, predicted = benign_outputs.max(1)

        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

    print('Train accuarcy:', 100. * correct / total)
    print('Train average loss:', train_loss / total)
    return (100. * correct / total, train_loss / total)


def evaluate(net, epoch, file_name, data_loader, info):
    print('[ Evaluate epoch: %d ]' % epoch)
    print("Dataset:", info)

    net.eval() # Dropout을 적용하는 경우 필수임
    test_loss = 0
    correct = 0
    total = 0

    for batch_idx, (inputs, targets) in enumerate(data_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        total += targets.size(0)

        outputs = net(inputs)
        test_loss += criterion(outputs, targets).item()

        _, predicted = outputs.max(1)
        correct += predicted.eq(targets).sum().item()

    print('Accuarcy:', 100. * correct / total)
    print('Average loss:', test_loss / total)
    return (100. * correct / total, test_loss / total)

 

3. Confusion Matrix 함수 정의

def get_confusion_matrix(net, num_classes, data_loader):
    confusion_matrix = torch.zeros(num_classes, num_classes)

    net.eval() # Dropout을 적용하는 경우 필수임

    for batch_idx, (inputs, targets) in enumerate(data_loader):
        inputs, targets = inputs.to(device), targets.to(device)

        outputs = net(inputs)
        _, predicted = outputs.max(1)

        for t, p in zip(targets.view(-1), predicted.view(-1)):
            confusion_matrix[t.long(), p.long()] += 1

    return confusion_matrix

 

3. ImageNet Pretrained ResNet을 활용한 Transfer Learning

net = torchvision.models.resnet18(pretrained=True)

# 마지막 레이어의 차원을 10차원으로 조절
num_features = net.fc.in_features
net.fc = nn.Linear(num_features, 10)
net = net.to(device)

epoch = 10
learning_rate = 0.01
file_name = "Transfer_Learning_MNIST.pt"

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0002)

train_result = []
test_result = []

train_result.append(evaluate(net, 0, file_name, transfer_train_loader, "Train"))
test_result.append(evaluate(net, 0, file_name, transfer_test_loader, "Test"))

for i in range(epoch):
    train(net, i, optimizer, criterion, transfer_train_loader)
    train_acc, train_loss = evaluate(net, i + 1, file_name, transfer_train_loader, "Train")
    test_acc, test_loss = evaluate(net, i + 1, file_name, transfer_test_loader, "Test")

    state = {
        'net': net.state_dict()
    }
    if not os.path.isdir('checkpoint'):
        os.mkdir('checkpoint')
    torch.save(state, './checkpoint/' + file_name)
    print('Model Saved!')

    train_result.append((train_acc, train_loss))
    test_result.append((test_acc, test_loss))

 

4. accurarcy 커브 시각화

# 정확도(accuracy) 커브 시각화
plt.plot([i for i in range(epoch + 1)], [i[0] for i in train_result])
plt.plot([i for i in range(epoch + 1)], [i[0] for i in test_result])
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend(["train", "test"])
plt.show()

 

5. loss 커브 시각화

# 손실(loss) 커브 시각화
plt.plot([i for i in range(epoch + 1)], [i[1] for i in train_result])
plt.plot([i for i in range(epoch + 1)], [i[1] for i in test_result])
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend(["train", "test"])
plt.show()

 

6. Training Dataset Confusion Matrix

# 혼동 행렬(Confusion Matrix) 시각화 (학습 데이터셋)
net = torchvision.models.resnet18(pretrained=True)

# 마지막 레이어의 차원을 10차원으로 조절
num_features = net.fc.in_features
net.fc = nn.Linear(num_features, 10)
net = net.to(device)

file_name = "./checkpoint/Transfer_Learning_MNIST.pt"
checkpoint = torch.load(file_name)
net.load_state_dict(checkpoint['net'])

confusion_matrix = get_confusion_matrix(net, 10, transfer_train_loader)
print("[ 각 클래스당 데이터 개수 ]")
print(confusion_matrix.sum(1))

print("[ 혼동 행렬(confusion matrix) 시각화 ]")
# 행(row)은 실제 레이블, 열(column)은 모델이 분류한 레이블
res = pd.DataFrame(confusion_matrix.numpy(), index = [i for i in range(10)], columns = [i for i in range(10)])
plt.figure(figsize = (10, 7))
sn.heatmap(res, annot=True)
plt.show()

print("[ 각 클래스에 따른 정확도 ]")
# (각 클래스마다 정답 개수 / 각 클래스마다 데이터의 개수)
print(confusion_matrix.diag() / confusion_matrix.sum(1))

print("[ 전체 평균 정확도 ]")
print(confusion_matrix.diag().sum() / confusion_matrix.sum())

 

7. Test Dataset Confusion Matrix

# 혼동 행렬(Confusion Matrix) 시각화 (테스트 데이터셋)
net = torchvision.models.resnet18(pretrained=True)

# 마지막 레이어의 차원을 10차원으로 조절
num_features = net.fc.in_features
net.fc = nn.Linear(num_features, 10)
net = net.to(device)

file_name = "./checkpoint/Transfer_Learning_MNIST.pt"
checkpoint = torch.load(file_name)
net.load_state_dict(checkpoint['net'])

confusion_matrix = get_confusion_matrix(net, 10, transfer_test_loader)
print("[ 각 클래스당 데이터 개수 ]")
print(confusion_matrix.sum(1))

print("[ 혼동 행렬(confusion matrix) 시각화 ]")
# 행(row)은 실제 레이블, 열(column)은 모델이 분류한 레이블
res = pd.DataFrame(confusion_matrix.numpy(), index = [i for i in range(10)], columns = [i for i in range(10)])
plt.figure(figsize = (10, 7))
sn.heatmap(res, annot=True)
plt.show()

print("[ 각 클래스에 따른 정확도 ]")
# (각 클래스마다 정답 개수 / 각 클래스마다 데이터의 개수)
print(confusion_matrix.diag() / confusion_matrix.sum(1))

print("[ 전체 평균 정확도 ]")
print(confusion_matrix.diag().sum() / confusion_matrix.sum())

8. 학습 이후 학습된 모델 다운로드 하기

from google.colab import files
files.download("이름.pt")

 

확실히 성능이 잘나오는 것을 확인할 수 있다!

Transfer Learning 짱!

 

 

전체코드와 결과는 깃헙에 있습니다

https://github.com/Kdavid2355/ai_code/blob/main/ImageNet_Pretrained_ResNet_for_MNIST_Classfication.ipynb