PyTorch 컴퓨터 비전 전이학습 예제를 따라해보자.

  이번 포스팅에서는 PyTorch 공식 페이지의 예제인 Transfer Learning for Computer Vision Tutorial을 따라해본다. 친절하게 설명된 양질의 글이 넘치는 세상에서 부족한 글을 쓰는 것은 항상 부끄럽지만, 더 좋은 글을 읽기 위해 훑어 볼 글 정도로 여겨졌으면 좋겠다.

  예제에서 설명되고 있는 전이학습(transfer learning)은 이미 대규모 데이터셋에서 훈련된 모델(pre-trained model)에서 출발하여 자신이 해결하려는 작업에 맞게 미세조정(fine-tuning)하여 훈련하는 과정을 말한다. 여기에서 사용한 ResNet-18 모델은 깊은 신경망을 훈련시킬 수 있는 스킵 연결(skip connection)을 특징으로 갖는 ResNet 모델의 하나로서, 18개의 컨볼루션 층을 가지고 있다. torchvision은 PyTorch에서 사용하는 컴퓨터 비전 라이브러리인데, torchvision에서 제공하는 resnet18은 1,000개의 클래스에 대한 1200만장의 ImageNet 데이터셋에서 사전 훈련된 모델을 제공한다. 남은 일은 개미 / 벌을 구분하도록 변형하여, 사용하는 것뿐이다.

 

필요한 라이브러리를 import 하자.

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

PyTorch를 사용하는 코드에서 가장 쉽게 볼 수 있는 import 문 형태는 위와 같을 것이다. 

• torch : 가장 기본이 되는 패키지로서, 다차원 텐서와 관련된 연산이 정의되어 있다.
• torch.nn : 신경망을 구성하는데 필요한 클래스와 모듈이 정의되어 있다.
• torch.nn.functional : torch.nn에 클래스로 정의된 모듈이 torch.nn.functional에 함수로 정의되어 사용될 수 있다.
• torch.optim : 다양한 최적화 알고리즘이 정의되어 있다.

 이번 예제에서는 추가적으로 컴퓨터 비전을 위한 torchvision과 학습 시 단계적으로 학습률을 조정하기 위한 torch.optim 패키지의 lr_scheduler를 사용한다. 그 외 시각화를 위한 matplotlib 등도 import 한다.

from torch.optim import lr_scheduler

import torchvision
from torchvision import datasets, models, transforms

import matplotlib.pyplot as plt
import numpy as np
import os
import time
import copy

데이터를 준비하자.

 예제에서 사용하는 데이터의 공개주소를 클릭하면, 개미와 벌의 이미지 데이터를 받을 수 있다. 

데이터 파일의 이름인 hymenoptera는 절지동물 곤충강 벌목의 한자명으로서, 벌과 개미가 여기에 속한다. 두 쌍의 막모양 날개를 가졌다.
솔직히 파일명을 내가 정했으면 ant_bee.zip이 최선이었을 것만 같다.

앞으로의 코드 진행을 위해서, 다운로드한 압축파일은 실행 중인 파이썬 파일이나 노트북 파일과 같은 디렉토리 안에 새로운 data 디렉토리를 만들고 압축을 풀어준다. 최종적으로 data/hymenoptera_data 안에 train 디렉토리와 val 디렉토리가 있으면 된다.

 

만약 구글 코랩(colab)을 이용해서 테스트를 진행 중인 경우라면, 사용 중인 세션의 현재 경로에 data 디렉토리를 만들어 진행하면 된다. 아래의 코드를 참고하면, 코랩에서 예제의 나머지 그 부분을 그대로 진행할 수 있다.

!wget --no-check-certificate https://download.pytorch.org/tutorial/hymenoptera_data.zip

import zipfile

zip_file = 'hymenoptera_data.zip'
zip_ref = zipfile.ZipFile(zip_file, 'r')
zip_ref.extractall('./data')
zip_ref.close()

!rm hymenoptera_data.zip

데이터를 로드하자.

 데이터가 준비되었으니 이제 할 일은 데이터를 메모리에 올리는 일이다. 컴퓨터가 데이터를 사용하기 위해서는 디스크에 저장된 데이터를 메인 메모리(main-memory)로 로드(load)하는 과정이 필요하다. 보통은 현대과학기술의 승리, '나는 믿을거야, 메인 메모리 믿을거야'를 부르짖으며 신경 쓰지 않고 싶어하지만, GPU를 사용하는 딥러닝 실험에서는 단순히 실행이 느려지는 정도에 그치지 않는다. 

main memory의 데이터가 GPU의 device memory로 전송될 때, device memory보다 더 큰 경우에는 메모리 할당 오류(cudaMalloc failed: out of memory)가 발생한다. Unified memory를 쓰면 어떻게든 되겠지만, 이런 걸 신경쓰지 않으려고 PyTorch 프레임워크를 사용하는 것 아니겠는가.

 컴퓨터 비전 분야의 대용량 데이터를 처리하기 위해서, 고맙게도 Pytorch에서는 DataLoader 클래스를 제공한다. 이를 사용하면 자동으로 배치를 생성(automatic batching)해주는 동시에 데이터 로딩 과정을 병렬화(parallelizing)할 수 있다. 단 하나, PyTorch의 데이터로더를 사용하기 위해서는 먼저 데이터셋을 만들어야 한다는 점만 알면 현재로서는 충분하다.

 DataLoader가 지원하는 dataset이 갖춰야할 요건은 특정 인덱스, 즉 순서에 맞는 이미지와 그 레이블을 짝지어서 제공해줘야 한다는 점이다. 실제 작업에서는 프로젝트 목적에 맞게 커스텀 데이터셋(custom dataset)을 만들어 쓰기도 하는데, 여기에서도 핵심은 여러 샘플을 순서에 맞게 인덱싱해주는 __getitem__ 메소드를 구현하는 것이다. (참고: Writing Custom Datasets, Dataloaders and Transforms)

 이번 예제에 사용하는 데이터는 PyTorch에서 전형적으로 사용하는 폴더 구조를 갖추고 있다. 즉, 각 클래스별로 이미지들이 서로 다른 폴더에 저장되어 있고, 클래스 폴더들도 하나의 폴더 안에 들어있다.

data/ hymenoptera_data / train / ants / xxxxx0.jpg
data/ hymenoptera_data / train / ants / xxxxx1.jpg
                                 …
data/ hymenoptera_data / train / bees / xxxxx0.jpg
data/ hymenoptera_data / train / bees / xxxxx1.jpg
                                 …

data/ hymenoptera_data / val / ants / xxxxx0.jpg
data/ hymenoptera_data / val / ants / xxxxx1.jpg
                                 …
data/ hymenoptera_data / val / bees / xxxxx0.jpg
data/ hymenoptera_data / val / bees / xxxxx1.jpg
                                 …

 이런 폴더 구조를 갖췄을 때, 손쉽게 dataset을 생성하는 함수를 torchvision 라이브러리는 제공한다. torchvision.datasets.ImageFolder를 호출하면 갖춰진 데이터 폴더 구조로부터 PyTorch dataset을 만들어준다.

 내부적으로는 앞서 언급했던 것처럼 샘플에 대한 인덱싱을 제공하고, 문자열로 된 클래스명들을 0번 클래스, 1번 클래스와 같이 정수 인덱스로 바꿔주는 작업이 이뤄지지만, 이런 복잡한 내부 작업을 숨기고 간단한 메소드 호출만으로 작업이 이뤄지게 하는 것이 라이브러리의 힘이다. 

ImageFolder 메소드는 ants / bees가 속한 train 폴더로부터 train dataset을, ants / bees가 속한 val 폴더로부터 validation dataset을 자동으로 만들어줄 것이다. 참고로 예제의 설명과 조금 다르게, train 폴더에 이미지 수는 (ants : bees ) = 124 : 121 이고, val 폴더에 이미지 수는 (ants : bees) = 70 : 83 장이다. 

 PyTorch dataset을 만들면서 torchvision 라이브러리를 쓰는 방식의 장점은 이미지 변형(image transformation)을 손쉽게 할 수 있다는 점이다. torchvision.transforms을 이용해서 변형 방식을 구성한 다음에 dataset을 만들 때 넘겨주는 것으로 충분하다. 이번 예제에 사용되는 transforms의 함수들을 정리하면 아래와 같다.

• transforms.Compose : 여러 변형 방식을 한꺼번에 묶어준다.
• transforms.RandomResizedCrop : 원본 이미지를 지정된 스케일 범위와 지정된 비율로 무작위로 자른 후에, 지정된 크기로 리사이즈 한다. 여기서는 scale과 ratio가 지정되지 않아서, 디폴트 값대로 원본 대비 0.08 ~ 1.0 사이의 크기와 가로/세로 비율 3/4 ~ 4/3 사이에서 크롭되고, 지정된 크기인 224에 따라서 (224, 224) 크기의 정사각 이미지로 리사이즈 된다.
• transforms.Resize : 입력 이미지를 주어진 크기로 리사이즈한다.
• transforms.CenterCrop : 입력 이미지의 중앙에서 시작하여 주어진 크기만큼 크롭한다. 참고로 입력 이미지 크기가 출력하려는 크기보다 작으면 부족한 부분의 픽셀값들이 0으로 패딩, 즉 채워진 다음에 크롭된다.
• transforms.HorizontalFlip : 주어진 확률로 이미지를 수평 반전 시킨다. 여기서는 확률값을 지정하지 않았으므로 디폴트 값인 0.5의 확률로 이미지들이 수평 반전 된다. 즉, 훈련 이미지 중 절반은 그대로지만, 절반은 뒤집히게 될 것이다.
• transforms.ToTensor : 우리가 사용한 ImageFolder 메소드를 비롯해서 torchvision 메소드는 이미지를 읽을 때, 파이썬 이미지 라이브러리인 PIL을 사용한다. 혹은 커스텀 데이터셋을 만들면서 PIL 대신에 OpenCV를 사용해서 이미지를 읽을 수도 있다. PIL을 사용해서 이미지를 읽으면 생성되는 Image 객체 또는 Image 객체로부터 변환되거나 OpenCV로 읽었을 때 생성된 numpy 배열은 픽셀값의 범위가 [0, 255] 이며, 배열의 차원이 (높이H x 너비W x 채널수C)로 되어 있다.
  그런데 이를 사전훈련된 모델 사용과 효율적 연산을 위해서 torch.FloatTensor 배열로 바꿔줘야 한다. 이 때, 픽셀값의 범위가 [0.0, 1.0] 사이가 되도록 바꿔줘야하며 차원의 순서를 바꿔서 (채널수C x 높이H x 너비W)가 되게 해야한다. 이 작업을 수행해주는 메소드가 ToTensor()다.
• transforms.Normalize : transfer learning에서 사용하는 pretrain model 들은 대개 ImageNet 데이터셋에서 사전훈련되어있다. 이 모델을 사용하기 위해서 ImageNet 데이터의 각 채널별 평균과 표준편차를 이용해 정규화(normalize)해 준다. 즉, Noramlize 메소드 안에 사용된 (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)는 ImageNet의 이미지들의 RGB 채널마다의 평균과 표준편차를 뜻한다. 참고로 OpenCV를 사용해 이미지를 읽어온다면 RGB이미지가 아닌 BGR이미지이므로, 채널 순서를 염두에 둬야 할 것이다.

 여기까지의 내용을 간단하게 요약하자면, 지정한 transform 방법을 가지고 dataset을 생성한 후에 이로부터 DataLoader를 만들어 데이터를 로드한다는 것이다. train dataset과 validation dataset에 대하여 각각 데이터셋, 데이터로더를 다른 변수명으로 생성하는 코드도 많지만( train_ds → train_dl, val_ds → val_dl ), 예제에서는 데이터셋과 데이터로더를 하나의 dictionary로 생성한 뒤에 key 값으로 구분하는 방식으로 생성하고 있다.

data_transform = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFilp(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'data/hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                        data_transforms[x]) for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(
                       image_datasets[x], batch_size=4, 
                       shuffle=True, num_workers=4)
                                            for x in ['train', 'val']}

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

모델을 학습하자.

 예제의 순서를 조금 바꿔서 사용하려는 모델을 먼저 로드하고, 손실함수와 옵티마이저 등을 먼저 설정하려고 한다. 처음에 언급한 것처럼 torchvision에서 제공하는 서브 패키지인 torchvision.models를 통해 ImageNet에서 사전 훈련된 ResNet 18 모델을 구성하려고 한다.

model_ft = models.resnet18(pretrained=True)

 이미 말한 것처럼 프레임워크를 쓰면 복잡한 내부 구현을 신경쓰지 않아도 된다. 그러나 조금만 더 들어가 보기로 한다. 이번 예제를 통해 한번만 알아두면 좋을 내용들이 있다. torchvision.models에 구현된 ResNet 모델의 소스 코드를 살펴본다. 

• ResNet 구현 코드 : torchvision / models / resnet.py

  긴 코드를 간단히만 말하자면 144번째 라인부터 시작되는 ResNet 클래스의 정의를 각 세부 모델의 구조에 맞게 252번째 라인의 _resnet 함수가 적절히 구성함으로써 모델이 만들어진다. ImageNet에서의 사전훈련을 고려해서인지 ResNet 클래스에서 디폴트로 지정된 클래스의 개수(num_classes)는 1000개로 되어 있다. ResNet 모델의 마지막 층인 fc 레이어는 기본적으로 num_classes에 해당하는 1000개의 노드를 가지고 있는 것이다. 

186:	self.fc = nn.Linear(512 * block.expansion, num_classes)

 이제 이 마지막 레이어를 우리가 해결하려는 문제에 맞게, 새로운 클래스 개수를 갖는 Linear 레이어를 만들어 대체하려고 한다. 새롭게 만드는 nn.Linear 레이어의 입력 개수(=num_ftrs)는 기존 Linear 레이어인 fc 레이어의 in_feature 속성값을 활용해서 알아내고, 출력 개수는 우리가 분류하려는 (개미 / 벌) 문제에 맞게 num_classes = 2 로 지정하면 된다.

num_ftrs = model_ft.fc.in_features

model_ft.fc = nn.Linear(num_ftrs, 2)

GPU 사용이 가능하다면 모델의 연산이 GPU에서 병렬적으로 수행될 수 있도록 장치에 할당한다.

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

model_ft = model_ft.to(device)

 

 그런데 여기서 생각해봐야 할 점은 torchvision.models 서브 패키지에서 정의된 모델의 마지막 층이 클래스별 확률분포를 알 수 있는 Softmax 층이 아니라 Linear 레이어라는 점이다.

 PyTorch에서는 마지막 출력 층에 Softmax를 사용하고, cross-entropy loss function을 사용하는 대신에 nn.LogSoftmax와 nn.NLLLoss가 결합된 nn.CrossEntropyLoss를 사용해서 분류 문제를 학습한다. PyTorch 공식문서에 따르면 CrossEntropyLoss를 사용하면 특히 클래스마다 데이터 수가 다른 분균형 학습셋(imbalanced traning set)의 학습에 유용하다고 한다.

Softmax 대신 사용하는 LogSoftmax에 대해 이해하려면 Yann LeCun 딥러닝 강의 2주차 자료를 살펴보면 좋을 것 같다.

 이에 따라 자연스럽게 목적함수(objective function = criterion)으로 CrossEntropyLoss를 활용할 수 있다.

criterion = nn.CrossEntropyLoss()

 마지막으로 결정해야 하는 것은 최적화를 위한 연산방식을 결정할 옵티마이저(Optimizer)를 지정하는 것이다. torch.optim 패키지는 쉽게 사용할 수 있는 여러 옵티마이저 구현체를 제공하는데, 이번 예제에서는 확률적 경사하강법(Stochastic Gradient Descent)을 실행하는 optim.SGD()를 사용하였다. 상세한 내용은 알아보지 않으려하는데, 앞으로 여러가지 옵티마이저를 이해해보기 위해서는 모멘텀(momentum)의 개념 등은 찾아보는 편이 더 좋을 것 같다.

optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)