[Machine Learning] PyTorch로 Transfer Learning (전이 학습) 구현하기

최신, 최첨단의 심층 네트워크를 만들어 써먹으려면 수일 혹은 수주일이 걸릴 수 있으며 수백만 개의 데이터가 필요하다.

예를 들어, ILSVRC ImageNet 2010에는 1000개의 카테고리가 있으며 각 카테고리마다 1000개가 넘는 고해상도 테스트 이미지도 있다.

모든 사람들이 처음부터 모델을 만들어 위와 같은 데이터를 분류하는 네트워크를 잘 훈련시키기에는 수많은 데이터가 필요하다.

이에 더해 최신 장비도 필요하고 시간도 많이 필요한 등 새로운 네트워크를 처음부터 훈련시키는 것은 매우 고된 작업이 될 것이다.

이런 에로사항 때문에 Transfer Learning(전이 학습)이 등장한다.

---

Transfer Learning

Transfer Learning이란 방대한 양의 데이터를 가지고 미리 잘 훈련된 모델을 가져와서 우리가 사용할 데이터에 맞게 모델을 수정한 다음 간단한 훈련을 통해 우리가 가진 데이터에 맞는 네트워크를 만드는 것이다.

주의할 것이 있는데 당연한 소리지만 우리가 훈련시킬 새로운 데이터들은 사전에 훈련된 네트워크가 학습했던 데이터와 연관이 있어야 한다.

예를 들어, 개의 품종을 분류하는 모델을 만들 때 범용적인 이미지 탐지기를 우리에 입맛에 맞게 수정하는 것은 가능하다.

하지만 위와 같은 이미지 네트워크로 바둑을 두는 역할을 맡기는 것은 좀 이상하고 아마 처음부터 훈련해야할 것이다.

 

이제 ILSVRC 데이터셋으로 훈련된 네트워크를 살펴보자.

Filters

다양한 레이어가 있는데 low-level 필터는 일반적인 다양한 각도의 선이나 곡선을 감지하는데 유용하고 mid-level 필터는 대부분의 작업에 유용해이고 high-level 필터는 광범위한 실제 이미지에 적합할 것이다.

 

따라서 위의 사전 훈련된 네트워크를 가져온 다음 우리가 트레이닝할 데이터 세트로 네트워크를 재교육한다면 처음부터 네트워크를 훈련시키는 것보다 더 적은 컴퓨팅 파워로 훈련이 가능하고 더 적은 시간만이 걸릴 것이다.

---

PyTorch로 Transfer Learning 구현하기

파이토치로 전이학습 코드를 직접 만들어보자.

vgg11 모델을 가져오는데 pretrained파라미터를 True로 설정하면 미리 사전에 훈련된 네트워크를 가져올 수 있다.

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

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

아래처럼 모델의 구성을 설명해준다.

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (4): ReLU(inplace=True)
    (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (6): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (7): ReLU(inplace=True)
    (8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (9): ReLU(inplace=True)
    (10): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (11): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (12): ReLU(inplace=True)
    (13): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (14): ReLU(inplace=True)
    (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (16): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (17): ReLU(inplace=True)
    (18): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (19): ReLU(inplace=True)
    (20): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

참고로 torchinfo을 사용하면 각 레이어의 아웃풋 사이즈를 볼 수 있다.

사용 방법은 이 글을 참고하자.

[PyTorch] torchinfo, summary 사용 방법 및 파라미터 / 여러 input 설정

[PyTorch] torchinfo, summary 사용 방법 및 파라미터 / 여러 input 설정

간단히 모델과 인풋 사이즈를 주면 아웃풋 사이즈를 바로 확인할 수 있다.

summary(net, (1, 3, 64, 64)) # model, (Batchsize, channels, width, height)

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
VGG                                      [1, 1000]                 --
├─Sequential: 1-1                        [1, 512, 2, 2]            --
│    └─Conv2d: 2-1                       [1, 64, 64, 64]           1,792
│    └─ReLU: 2-2                         [1, 64, 64, 64]           --
│    └─MaxPool2d: 2-3                    [1, 64, 32, 32]           --
│    └─Conv2d: 2-4                       [1, 128, 32, 32]          73,856
│    └─ReLU: 2-5                         [1, 128, 32, 32]          --
│    └─MaxPool2d: 2-6                    [1, 128, 16, 16]          --
│    └─Conv2d: 2-7                       [1, 256, 16, 16]          295,168
│    └─ReLU: 2-8                         [1, 256, 16, 16]          --
│    └─Conv2d: 2-9                       [1, 256, 16, 16]          590,080
│    └─ReLU: 2-10                        [1, 256, 16, 16]          --
│    └─MaxPool2d: 2-11                   [1, 256, 8, 8]            --
│    └─Conv2d: 2-12                      [1, 512, 8, 8]            1,180,160
│    └─ReLU: 2-13                        [1, 512, 8, 8]            --
│    └─Conv2d: 2-14                      [1, 512, 8, 8]            2,359,808
│    └─ReLU: 2-15                        [1, 512, 8, 8]            --
│    └─MaxPool2d: 2-16                   [1, 512, 4, 4]            --
│    └─Conv2d: 2-17                      [1, 512, 4, 4]            2,359,808
│    └─ReLU: 2-18                        [1, 512, 4, 4]            --
│    └─Conv2d: 2-19                      [1, 512, 4, 4]            2,359,808
│    └─ReLU: 2-20                        [1, 512, 4, 4]            --
│    └─MaxPool2d: 2-21                   [1, 512, 2, 2]            --
├─AdaptiveAvgPool2d: 1-2                 [1, 512, 7, 7]            --
├─Sequential: 1-3                        [1, 1000]                 --
│    └─Linear: 2-22                      [1, 4096]                 102,764,544
│    └─ReLU: 2-23                        [1, 4096]                 --
│    └─Dropout: 2-24                     [1, 4096]                 --
│    └─Linear: 2-25                      [1, 4096]                 16,781,312
│    └─ReLU: 2-26                        [1, 4096]                 --
│    └─Dropout: 2-27                     [1, 4096]                 --
│    └─Linear: 2-28                      [1, 1000]                 4,097,000
==========================================================================================
Total params: 132,863,336
Trainable params: 132,863,336
Non-trainable params: 0
Total mult-adds (M): 735.31
==========================================================================================
Input size (MB): 0.05
Forward/backward pass size (MB): 4.92
Params size (MB): 531.45
Estimated Total Size (MB): 536.43
==========================================================================================

아무튼 여기서 마지막 레이어는 보면 4096개의 입력을 받아서 1000개의 클래스로 구분을 하는 레이어임을 알 수 있다.

그리고 학습 가능한 파라미터가 1억 3천만개가 넘는 거대한 네트워크임을 알 수 있다.

 

이 네트워크를 가지고 10개 클래스의 옷이나 신발을 구분하는 FashionMNIST 네트워크를 만들어보자.

네트워크의 마지막 레이어인 classifier[6]의 아웃풋을 1000에서 10으로 변경하면 된다.

하지만 먼저 해야할 중요한 것은 사전 학습된 파라미터는 일단 건드리지 않는 것이 전이학습의 아이디어이다.

따라서 마지막 레이어가 아닌 레이어들의 가중치들은 그래디언트를 추적하지 않게 만들자.

for p in net.parameters():
    p.requires_grad = False

이따 새롭게 추가할 레이어와 비교하기 위해 아래 코드를 실행시켜보자.

파라미터의 requires_grad가 True인지 False인지를 출력해준다.

for p in net.parameters():
    print("requireds_grad:", p.requires_grad)

위 루프의 실행 결과는 모두 required_grad: False가 나올 것이다.

즉, 현재 있는 모든 파라미터(사전 학습된 가중치)의 역전파를 계산하지 않는다는 소리다.

 

이제 학습에 사용할 데이터를 가져오자.

trans = transforms.Compose([
    transforms.Resize(224), # upscale from 28x28 to 224x224
    transforms.Grayscale(3), # turns it into a 3 color image with channels r==g==b
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train = torchvision.datasets.FashionMNIST('FashionMNIST', train=True, download=True, transform=trans)
test = torchvision.datasets.FashionMNIST('FashionMNIST', train=False, download=True, transform=trans)
trainloader = torch.utils.data.DataLoader(train, batch_size=64, num_workers=4)
testloader = torch.utils.data.DataLoader(test, batch_size=64, num_workers=4)

이번엔 transform 파라미터에 대해 뭐가 좀 많은데, vgg11 모델이 트레이닝 되었던 환경으로 데이터도 그에 맞게 변형해준 것이다.

자세한 점은 PyTorch의 vgg11에 대한 문서에서 확인할 수 있다.

 

이제 마지막 레이어의 아웃풋을 1000에서 10으로 바꾸자.

새롭게 추가되는 레이어는 기본적으로 requires_grad가 True이다.

net.classifier[6] = nn.Linear(4096, 10)

진짜 그런지 아래 코드를 돌려보면 알 수 있다.

for p in net.parameters():
    print("requireds_grad:", p.requires_grad)

새롭게 바뀐 네트워크도 summary(net, (1, 3, 64, 64))을 돌려서 마지막 레이어가 바뀌었는지 확인해보자.

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
VGG                                      [1, 10]                   --
├─Sequential: 1-1                        [1, 512, 2, 2]            --
│    └─Conv2d: 2-1                       [1, 64, 64, 64]           (1,792)
│    └─ReLU: 2-2                         [1, 64, 64, 64]           --
│    └─MaxPool2d: 2-3                    [1, 64, 32, 32]           --
│    └─Conv2d: 2-4                       [1, 128, 32, 32]          (73,856)
│    └─ReLU: 2-5                         [1, 128, 32, 32]          --
│    └─MaxPool2d: 2-6                    [1, 128, 16, 16]          --
│    └─Conv2d: 2-7                       [1, 256, 16, 16]          (295,168)
│    └─ReLU: 2-8                         [1, 256, 16, 16]          --
│    └─Conv2d: 2-9                       [1, 256, 16, 16]          (590,080)
│    └─ReLU: 2-10                        [1, 256, 16, 16]          --
│    └─MaxPool2d: 2-11                   [1, 256, 8, 8]            --
│    └─Conv2d: 2-12                      [1, 512, 8, 8]            (1,180,160)
│    └─ReLU: 2-13                        [1, 512, 8, 8]            --
│    └─Conv2d: 2-14                      [1, 512, 8, 8]            (2,359,808)
│    └─ReLU: 2-15                        [1, 512, 8, 8]            --
│    └─MaxPool2d: 2-16                   [1, 512, 4, 4]            --
│    └─Conv2d: 2-17                      [1, 512, 4, 4]            (2,359,808)
│    └─ReLU: 2-18                        [1, 512, 4, 4]            --
│    └─Conv2d: 2-19                      [1, 512, 4, 4]            (2,359,808)
│    └─ReLU: 2-20                        [1, 512, 4, 4]            --
│    └─MaxPool2d: 2-21                   [1, 512, 2, 2]            --
├─AdaptiveAvgPool2d: 1-2                 [1, 512, 7, 7]            --
├─Sequential: 1-3                        [1, 10]                   --
│    └─Linear: 2-22                      [1, 4096]                 (102,764,544)
│    └─ReLU: 2-23                        [1, 4096]                 --
│    └─Dropout: 2-24                     [1, 4096]                 --
│    └─Linear: 2-25                      [1, 4096]                 (16,781,312)
│    └─ReLU: 2-26                        [1, 4096]                 --
│    └─Dropout: 2-27                     [1, 4096]                 --
│    └─Linear: 2-28                      [1, 10]                   40,970
==========================================================================================
Total params: 128,807,306
Trainable params: 40,970
Non-trainable params: 128,766,336
Total mult-adds (M): 731.25
==========================================================================================
Input size (MB): 0.05
Forward/backward pass size (MB): 4.92
Params size (MB): 515.23
Estimated Total Size (MB): 520.19
==========================================================================================

결과를 보면 마지막 레이어의 Output Shape이 [1, 10] # [임의로 입력한 배치 사이즈, 클래스 개수]임을 알 수 있다.

맨 오른쪽 줄의 Param을 보면 기존 레이어의 파라미터 숫자가 괄호 안에 있는 것도 확인할 수 있다.

이것은 우리가 아까 requires_grad = False로 설정했기 때문이다.

 

이제 학습할 파라미터들을 따로 리스트에 담아 옵티마이저에 담을 준비를 하자.

params = []
for p in net.parameters():
    if p.requires_grad:
        params.append(p)

옵티마이저와  크로스 엔트로피도 만들자.

optim = torch.optim.SGD(params, lr=1e-2, momentum=0.9)
criterion = torch.nn.CrossEntropyLoss()
device = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu') # For m1 Mac Users
net = net.to(device)

아무튼 이제 데이터도 준비가 됐고 사전 학습이 완료된 vgg11기반의 네트워크도 있고 네트워크의 말단도 우리가 원하는 대로 설정하였고 학습할 파라미터도 따로 준비가 됐다.

 

이제는 트레이닝 루프를 만들고 결과를 확인하는 이전 글에서 많이 해봤던 작업만이 남았다.

트레이닝 루프는 본인이 직접 만들어도 된다.

나는 그냥 테스트 데이터셋의 정확도만 모니터링하는 코드를 대충 만들어 보았다.

# Train the model
total_step = len(trainloader)
num_epochs=5
for epoch in range(1, num_epochs+1):
    for i, (images, labels) in enumerate(trainloader):
        net.train()
        images = images.to(device)
        labels = labels.to(device)
        
        # Forward pass
        outputs = net(images)
        loss = criterion(outputs, labels)
        
        # Backward and optimize
        optim.zero_grad()
        loss.backward()
        optim.step()
        
        if (i+1) % 100 == 0:
            net.eval()
            with torch.no_grad():
                correct, total = 0, 0
                for images, labels in testloader:
                    images = images.to(device)
                    labels = labels.to(device)
                    outputs = net(images)
                    _, predicted = torch.max(outputs.data, 1)
                    total += labels.size(0)
                    correct += (predicted == labels).sum().item()
            print ('Epoch [{}/{}], Step [{}/{}], Acc: {:.4f} %' 
                   .format(epoch, num_epochs, i+1, len(trainloader), 100*correct/total))

(루프를 다 도는데 한 시간 가까이 걸린다.)

---

학습할 파라미터들이 1억 3천만개가 넘는 이런 거대한 네트워크를 처음부터 학습시키는 것보단 이미 좋은 성능의 네트워크를 끌고와서 나머지 40,970개의 파라미터를 트레이닝 하는게 훨씬 편하다.

torchvision.models를 보면 매우 다양한 네트워크가 있다.

자세한 내용은 공식 문서를 통해 확인해보자.

확인해보고 자신이 트레이닝할 데이터와 비슷한 데이터로 훈련된 네트워크를 찾아 Transfer Learning을 통해 시간과 자원을 아껴보자.