Transfer Learning 을 통한 나만의 이미지 분류기 만들기

Heath_Jeong·2021년 3월 21일

뚝딱뚝딱

나만의 데이터셋으로 이미지 분류 모델 만들기

목록 보기

2/2

개요

이전 포스트를 통해 마동석, 수지, 유재석 데이터셋이 준비되었다.

이제 전이 학습 (Transfer Learning) 을 통해 마동석, 수지, 유재석 레이블을 구분하는 이미지 분류기를 만들어 보자.

절차
1) train_iter, test_iter 준비
2) pretrained 모델 불러오기 (ResNet18)
3) FC Layer 구조 변경
4) 학습
5) 테스트

train_iter, test_iter 준비

이전 포스트에서 sklearn 의 train_test_split 를 통해 데이터가 3:1 비율로 train 과 test 로 분리되었다.

먼저 datasets.ImageFolder 를 통해 분리된 train, test 폴더에서 데이터를 불러온다.

from torchvision import datasets, transforms, models

DATASET_PATH = "/content/drive/MyDrive/부스트캠프 AI Tech 4조/data_team4/동석수지재석/revised/dataset"

train_data = datasets.ImageFolder(DATASET_PATH + '/train', transform=train_transforms)
test_data = datasets.ImageFolder(DATASET_PATH + '/test', transform=test_transforms)

그리고 데이터를 부를 때, 모델이 데이터를 더 잘 이해할 수 있도록 transforms.Compose 함수로 Data Augmentation 을 수행한다.

train_transforms = transforms.Compose([transforms.RandomRotation(30),
                                       transforms.CenterCrop(224),
                                       transforms.ToTensor(),
                                       transforms.Normalize([0.485, 0.456, 0.406],
                                                            [0.229, 0.224, 0.225])])

test_transforms = transforms.Compose([transforms.Resize(255),
                                      transforms.CenterCrop(224),
                                      transforms.ToTensor(),
                                      transforms.Normalize([0.485, 0.456, 0.406],
                                                           [0.229, 0.224, 0.225])])

DataLoader 함수를 통해 모델에 데이터를 넣을 수 있도록 train_iter, test_iter 로 만든다.

train_iter = torch.utils.data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=1)
test_iter = torch.utils.data.DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=1)

pretrained 모델 불러오기 (ResNet18)

Pytorch 의 torchvision.models 을 통해 직접 모델을 정의하지 않고도 이미지 처리와 관련된 모델들을 사용할 수 있다.
심지어 pretrained=True 옵션을 통해 학습된 모델을 불러올 수도 있다.

import torchvision.models as models

pretrained_model = models.resnet18(pretrained=True)

FC Layer 구조 변경

그렇다면 학습된 모델을 불러와서 어떻게 원하는 분류 태스크를 수행할 수 있을까?

전이학습 (Transfer Learning)

전이학습은 많은 데이터로 미리 학습한 모델에 원하는 태스크를 수행하도록 새로운 데이터를 추가적으로 학습시키는 방법을 말한다.

CNN 의 구조를 살펴보면 크게 Convolution Layers 파트와 Fully Connected Layers 파트가 있다.

Convolution Layers 는 이미지를 이해하는 단계이다.
초반의 레이어에서는 이미지의 선이나 점 등 작은 특징을 구분해내고 뒷쪽의 레이어로 갈 수록 원, 네모 등 더 뚜렷하게 이미지의 특징을 이해한다.

Fully Connected Layers 는 이미지에 대한 이해를 바탕으로 전체 레이어를 하나로 펴서 레이블을 분류한다.

적용

전이학습을 위해 앞의 Convolution Layers 는 그대로 사용하고 Fully Connected Layers 를 수정하여 원하는 레이블을 분류하도록 학습할 것이다.

아래의 코드를 보면 분류하고자 하는 레이블이 3 개이기 때문에 모델의 fc 의 out_degree 를 3 으로 세팅하였다.

num_ftrs = pretrained_model.fc.in_features
pretrained_model.fc = nn.Linear(num_ftrs, 3)

학습

이미지를 입력으로 받으면 마동석, 수지, 유재석 를 구분하는 모델을 만들기 위해 학습을 진행한다.

for epoch in range(EPOCHS):
    loss_val_sum = 0
    for batch_img, batch_lab in train_iter:

        X = batch_img.to(device)
        Y = batch_lab.to(device)
        
        # Inference & Calculate loss 
        y_pred = pretrained_model.forward(X)
        loss = criterion(y_pred, Y)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        loss_val_sum += loss
        
    if ((epoch%print_every)==0) or (epoch==(EPOCHS-1)):
        # accr_val = M.test(x_test, y_test, batch_size)
        loss_val_avg = loss_val_sum / len(train_iter)
        accr_val = test_eval(pretrained_model, test_iter, BATCH_SIZE)
        print(f"epoch:[{epoch+1}/{EPOCHS}] cost:[{loss_val_avg:.3f}] test_accuracy:[{accr_val:.3f}]")

학습 결과

10 에포크 밖에 안되지만 cost (loss) 도 낮고, 정확도도 높은 것을 확인할 수 있다.
epoch:[10/10] cost:[0.127] test_accuracy:[95.486]

테스트

test_iter 를 통해 학습된 모델의 분류 성능을 확인한다.

n_sample = 16
# sample_indices = np.random.choice(len(mnist_test.targets), n_sample, replace=False)
test_x = images[:n_sample]
test_y = labels[:n_sample]

with torch.no_grad():
    pretrained_model.eval()
    y_pred = pretrained_model.forward(test_x.type(torch.float).to(device))
    pretrained_model.train()
    
y_pred = y_pred.argmax(axis=1)

plt.figure(figsize=(20, 20))

test_accuracy = sum([1 for i, j in zip(y_pred, test_y) if i == j]) / n_sample
print(f"test_accuracy : {test_accuracy:.3f}")

for idx in range(n_sample):
    ax = plt.subplot(4, 4, idx+1)
    title = f"Predict: {y_pred[idx]}, Label: {test_y[idx]}"
    imshow(test_x[idx], ax, title)

plt.show()