[딥러닝 기초] 전이 학습 (Transfer learning)

1. 전이 학습
2. ImageNet
3. 레이어 동결
4. 사전학습 모델 선택하기
5. 모델 저장과 불러오기
    5.1. PyTorch
    5.2. Tensorflow

1. 전이 학습 (Transfer learning)

성능 좋은 딥러닝 모델을 만드는 데는 많은 데이터와 그 데이터를 감당할 만한 큰 모델을 필요로 합니다. 그러나 대부분의  일반인들이 이 많은 데이터와 모델을 학습시키기에는 컴퓨터 성능이 부족할 겁니다. 설령 컴퓨터 성능이 받쳐준다고 할지라도, 그 많은 데이터를 학습 시키는데 시간이 너무 오래 걸립니다. 딥러닝 모델을 만들 때마다 이 만큼의 자원과 시간을 소모하는 것은 매우 힘든 일이겠죠.

그래서 이렇게 큰 시간을 들여서 미리 학습한 모델을 가져다 쓸 수는 없을까 하는 생각을 하게 됩니다. 이렇게 대량의 데이터로 긴 시간 미리 학습하는 것을 '사전 학습(Pre-training)'이라고 합니다. 그리고 사전학습된 모델을 새로운 학습에 활용하는 것을 '전이 학습(Transfer learning)'이라고 하며, 사전학습된 모델의 파라미터들을 새로운 task를 통해 업데이트 시키는 것을 '미세 조정(Fine tuning)'이라고 합니다. 엄연히 따지면 transfer learning이 fine tuning을 포괄하는 개념이지만 혼용해서 많이 사용됩니다.

사전 학습은 단순히 학습을 미리한다는 것 이상의 의미가 있습니다. 왜냐하면 사전 학습에 사용되지 않은 데이터에 대해서도 좋은 성능을 보이기 때문입니다. 사람으로 예를 들면 기타를 칠 줄 아는 사람이 베이스도 쉽게 배우는 것과 비슷하다고 볼 수 있습니다. 인공지능 모델 역시 유사한 데이터 학습을 통해 다른 데이터에 대해서도 좋은 성능을 낼 수 있게 되는 겁니다.

이를 통해 성능 좋은 인공지능 모델을 좀 더 많은 사람이 자신의 문제 해결에 활용할 수 있게 되었으니 이야말로 정말 훌륭한 기술이 아닌가 싶습니다.

2. ImageNet

가장 대표적인 사전 학습 모델은 ImageNet을 학습한 모델인 것 같습니다. ImageNet은 1400만 개 이상의 이미지로 구성되어 있으며, 1000개의 분류 라벨을 갖는 데이터셋입니다. 이미지들은 '동물', '식물', '가구', '악기' 등 일상에서 쉽게 접할 수 있는 주제의 단어들을 키워드로 하여 flickr나 다양한 검색 엔진에서 검색을 통해 수집했다고 합니다.

그렇기 때문에 일상의 이미지들을 이용하는 대부분의 일상 task들에는 이 ImageNet으로 학습한 모델이 좋은 성능을 보입니다. 그래서 Tensorflow나 PyTorch에서도 기본적으로 쉽게 모델을 불러올 수 있도록 코드를 제공하고 있습니다. 이 기본 모델들의 크기나 성능에 대해서 keras에 표로 잘 정리되어 있습니다.

 

일단 이번엔 좀 옛날 모델이긴 하지만, 사전학습 모델의 대명사 격인 VGG16을 사용해보겠습니다. VGG16은 Convolution Block(Convolution layer와 Pooling layer의 집합)을 여러개 쌓은 구조를 갖고 있습니다.

VGG16 구조

모델은 다음 코드를 이용해 쉽게 불러올 수 있습니다.

from torchvision.models import vgg16

vgg = vgg16(weights="IMAGENET1K_V1")

 

지금부터 사전학습의 효과를 시험해 보기 위해 CIFAR-10 데이터를 이용해 실험을 진행해 볼 겁니다.

https://colab.research.google.com/drive/1oYTQ0tMxPfxPPGyVsm2CHmH51zt-Bae5?usp=sharing 

Transfer learning.ipynb

코드는 위 colab 링크에 공유해 두었습니다.

 

설명만 듣고 직접 전이학습을 구현해 보려고할 때 가장 먼저 부딪히게 될 문제는 아마도 모델의 크기가 안 맞는 문제일 것입니다. 저희가 불러온 VGG16은 224x224 크기의 이미지로 학습된 모델입니다. 그러나 CIFAR-10 데이터셋은 32x32 크기의 이미지로 구성되어 있습니다. 224x224 크기의 이미지가 Convolution Block들을 거친 뒤에 출력되는 크기는 (7x7x512)입니다. 하지만 32x32 크기의 이미지가 Convolution Block들을 거친 뒤에 출력되는 크기는 (1x1x512)입니다.

imagenet = torch.randn(1, 3, 224, 224) # (배치 크기, 이미지 채널 수, 가로 길이, 세로 길이)
orig_size = model.features(imagenet)

cifar = torch.randn(1, 3, 32, 32)
cifar_size = model.features(cifar)
print("ImageNet 데이터를 입력했을 때 출력되는 크기 :", orig_size.size())
print("Cifar-10 데이터를 입력했을 때 출력되는 크기 :", cifar_size.size())

'''
ImageNet 데이터를 입력했을 때 출력되는 크기 : torch.Size([1, 512, 7, 7])
Cifar-10 데이터를 입력했을 때 출력되는 크기 : torch.Size([1, 512, 1, 1])
'''

이것이 문제가 되는 것은 Convolution layer 뒤에 Linear layer가 연결되어 있기 때문입니다. Convolution layer는 이미지의 크기가 달라도 상관 없는 반면에 Lienar layer는 고정된 크기의 input을 요구합니다. 사전학습된 VGG16의 Linear 레이어는 512x7x7개의 feature를 input으로 요구하지만, CIFAR-10 이미지를 입력하면 512개의 feature가 나오기 때문에 shape이 맞지 않는 에러가 발생하게 되는 것입니다.

RuntimeError : mat1 and mat2 shapes cannot be multiplied (512x1 and 25088x4096)

그렇기 때문에 새로 학습하려는 데이터에 맞춰 모델의 구조를 살짝 바꿔줄 필요가 있습니다. 512개의 feature를 입력 받아 CIFAR-10의 분류 라벨 수인 10개의 output을 내는 새로운 linear 레이어를 VGG16의 Convolution block 뒤에 새로 붙여줄겁니다.

class VGGforCIFAR(nn.Module):
    def __init__(self, vgg):
        super(VGGforCIFAR, self).__init__()
        self.vgg_layers = vgg
        self.linear = nn.Linear(512, 10)

    def forward(self, inputs):
        feature = self.vgg_layers(inputs)
        output = self.linear(feature.squeeze())

        return output

vgg = VGGforCIFAR(model.features)

VGG16의 Convolution Block만 가져온 뒤, Linear 레이어는 CIFAR-10에 맞춰 새로 만든다.

사전학습의 효과를 확인하기 위해, 똑같은 모델 구조를 갖지만 사전학습된 weight를 사용한 모델과 사용하지 않은 모델을 둘 다 학습해 보고 성능을 비교해 보겠습니다. (데이터가 부족한 일상의 상황을 가정하기 위해 전체 5만 개의 이미지 중 5천개만 훈련 데이터로 사용했습니다.)

from torchvision.models import vgg16

model = vgg16(weights=None)
pretrained_model = vgg16(weights="IMAGENET1K_V1")

vgg = VGGforCIFAR(model.features)
pretrained_vgg = VGGforCIFAR(pretrained_model.features)

 

사전학습된 weight를 쓴 모델과 쓰지 않은 모델의 정확도 차이

결과 그래프를 확인해 보면, 같은 모델구조임에도 불구하고 사전학습된 weight를 사용한 모델의 성능이 2배 가까이 좋은 것을 확인할 수 있습니다. 이를 보면 사전학습이 모델의 성능에 정말 큰 영향을 미친다는 것을 알 수 있습니다.

3. 레이어 동결

여기서 이런 생각이 들 수 있습니다. "VGG16의 convolution block은 사전학습된 weight를 사용하는데 완전 랜덤하게 초기화된 weight들은 추가해서 학습하면 조금 이상하지 않을까?" 실제로 생각해 보겠습니다. 사전학습된 VGG16 conv block에 학습이 하나도 안된 linear 레이어를 처음부터 함께 학습한다면 어떨까요?

새로운 데이터로 학습이 하나도 되지 않았기 때문에 맨 초기의 모델 상태는 위와 같이 optimal 지점과 멀 것입니다. 이 상태에서 역전파를 수행한다면 미분값이 커지면서 모델의 파라미터가 크게 변화하게 됩니다. 이 때 VGG16의 convolution block도 함께 학습한다면, 이들의 파라미터도 크게 변화하게 되겠죠. 즉, 원래 사전학습 되었던 정보들을 잃기 쉬워 모델 성능에 악영향을 줄 수 있습니다.

 

그렇기 때문에 기존의 레이어들은 학습에서 '제외'해 줄 필요가 있습니다. 모델 레이어의 가중치를 업데이트 하지 않는 것을 해당 레이어를 '동결(freeze)'한다고 합니다. 레이어의 동결은 아래와 같이 trainable=False를 통해 설정할 수 있습니다.

vgg.features.trainable = False

 

 

일반적으로 전이학습은 아래와 같은 단계를 거쳐서 학습합니다.

1. 사전학습된 레이어를 동결하고 새로운 레이어만 파라미터 업데이트를 진행하며 학습을 수행.
2. 학습이 안정화되면 사전학습된 레이어의 동결을 해제.
3. 사전학습된 레이어의 정보를 최대한 유지하기 위해 learning rate를 줄여가며 추가 학습.

위 단계에 따라 CIFAR-10 데이터를 다시 학습해 보겠습니다. Linear 레이어가 어느정도 안정화 되도록 VGG16 레이어를 동결하고 1에포크 학습한 뒤, 동결을 해제한 뒤에 learning rate를 낮춰 다시 5에포크 추가 학습을 수행해 봤습니다.

좀 더 세밀하게 fine-tuning을 수행한 결과(초록색) 정확도를 더 높일 수 있었다.

결과는 위와 같이 모델 전체를 학습한 것보다 더 높은 정확도를 달성할 수 있었습니다. 이렇게 fine-tuning을 얼마나 세밀하게 수행하느냐에 따라서도 성능 차이가 발생하게 됩니다.

4. 사전학습 모델 선택하기

좋은 성능의 모델을 얻기 위해서 사전 학습 모델을 사용하는 것은 거의 필수라고 볼 수 있습니다. 그리고 감사하게도 많은 사람들이 사전학습 모델을 공개하고 있습니다. 사전 학습 모델을 찾는 방법은 여러가지가 있습니다.

 

1. Github에서 직접 찾기 : Github에 코드와 사전학습 모델을 함께 공개하는 경우가 있습니다. 이 repo를 직접 다운 받아 사용하는 방법이 있습니다.
2. PyTorch hub, Tensorflow hub, Huggingface와 같은 모델 공유 커뮤니티 활용 : 사전학습 모델들을 쉽게 공유하고 다운 받을 수 있도록 해주는 플랫폼입니다.

이런 여러가지 방법을 활용하여 자신의 task에 알맞는 사전학습 모델을 사용하면 됩니다. 적절한 사전학습 모델을 찾기 위해 고려해야 하는 사항이 몇가지 있습니다.

가장 먼저 고려해야 하는, 가장 중요한 사항은 사전학습 task가 내 task와 얼마나 유사한가 입니다. 이는 사전학습 방식을 말하는 거기도 하고, 사전학습에 사용된 데이터를 말하기도 합니다.

아무래도 수행하고자 하는 학습 방법과 동일한 방식으로 사전학습된 모델이 더 좋은 성능을 보일 겁니다. 데이터도 마찬가지입니다. 사전학습에 사용된 데이터가 제가 준비한 데이터의 특성과 유사할수록 사전학습된 모델이 새로운 데이터에 대해서도 더 잘 적응할 겁니다.

다음으로 생각해 볼 수 있는 것은 성능일 겁니다. 이러니 저러니 해도 제가 하려는 task에서 가장 성능이 좋은 사전학습 모델을 쓰는 것이 가장 좋겠죠. 컴퓨터가 정말 좋다면 가장 성능 좋은 모델을 불러와서 사용하면 될겁니다. 하지만 보통 가장 성능 좋은 모델은 정말 거대하고 여러분의 컴퓨터는 그렇지 않겠죠... 그렇기 때문에 모델의 성능과 내 컴퓨터가 감당 가능한 모델의 크기를 함께 고려하여 적당히 타협하여 모델을 선택해야 합니다.

5. 모델 저장과 불러오기

사전학습을 직접 수행하고 모델을 저장할 수도 있을 겁니다. 그리고 이 모델들을 불러와야 할 때도 있겠죠. PyTorch와 Tensorflow 모두 모델 저장과 불러오기 기능을 당연히 제공하고 있습니다.

5.1. PyTorch

PyTorch는 모델의 파라미터들을 dictionary 형태로 저장할 수 있습니다.

import torch

PATH = "./models/cifar_model.pt"
torch.save(model.state_dict(), PATH)

state_dict()는 아래와 같이 레이어의 이름과 파라미터 값들을 저장하고 있습니다.

{'features.0.weight': tensor([[[[-0.0161, -0.0084, 0.0495], [-0.0206, -0.0102, 0.0864], [ 0.0016, -0.0016, 0.0258]], ....

이렇게 저장한 파라미터들은 이 파라미터를 사용한 모델과 똑같은 모델을 구현한 뒤, 파라미터를 덮어씌우는 방식으로 사용할 수 있습니다.

from torchvision.models import vgg16

vgg = vgg16(weights="IMAGENET1K_V1")

class VGGforCIFAR(nn.Module):
    def __init__(self, vgg):
        super(VGGforCIFAR, self).__init__()
        self.vgg_layers = vgg
        self.linear = nn.Linear(512, 10)

    def forward(self, inputs):
        feature = self.vgg_layers(inputs)
        output = self.linear(feature.squeeze())

        return output

model = VGGforCIFAR(vgg.features)
# 파라미터 불러오기
pretrained_weights = torch.load("./models/cifar_model.pt")
model.load_state_dict(pretrained_weights)

모델을 다시 구현하는 것이 귀찮다면 모델 자체를 저장할 수도 있습니다.

torch.save(model, "models/cifar_model.pth")

model = torch.load("models/cifar_model.pth")

PyTorch는 모델 파일의 확장자로 ".pt", ".pth", ".bin"을 사용합니다.

5.2. Tensorflow

Tensorflow도 PyTorch와 같이 모델 파라미터를 저장하고 불러올 수 있습니다.

def create_model():
    model = tf.keras.models.Sequential()
    model.add(tf.keras.layers.Dense())
    ...
    return model

model = create_model()
model.save_weights('./checkpoints/model1')

new_model = create_model()
new_model.load_weights('./checkpoints/model1')

 

다만 tensorflow의 경우 폴더 형태로 저장하고, 불러올 때도 폴더를 불러오는 식으로 합니다.

혹은 모델 전체를 저장할 수도 있습니다. 역시 폴더 채로 저장합니다.

model.save("./model")

new_model = tf.keras.models.load_model("./model")

폴더가 불편하면 하나의 파일로 저장할 수도 있습니다. Tensorflow는 '.h5' 확장자를 사용합니다.

model.save("example.h5")

new_model = tf.keras.models.load_model("example.h5")