[파이토치] Transfer Learning

전이학습(Transfer Learning)

전이학습이란 기존의 잘 알려진 데이터 혹은 사전학습 된 모델을 업무 효율 증대나 도메인 확장을 위해 사용하는 학습을 의미한다. 모델의 초기 low level layer는 general한 feature를 추출하고 high level에서는 specific한 특징을 추출해내는 고도화된 학습이 이루어진다. 따라서 초기 layer의 feature들은 학습할 때 재사용해도 되지만 후반부에 위치한 layer의 feature들은 재학습이 필요하다. 본 실습은 CIFAR-10 데이터를 이용하였다.

10개의 클래스로 이루어진 CIFAR-10 dataset

 

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import torch.nn as nn
import torch.optim as optim
# 데이터 불러오기 및 전처리 작업
transform = transforms.Compose(
    [transforms.RandomCrop(32, padding=4), # data augmentation 기법: 이미지 한장에서 일부를 crop한다.
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

test_transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True) 

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=16,shuffle=False)

 

 

다음은 ResNet18 불러오겠습니다. pretrained=True를 하면 ResNet18 구조와 사전 학습 된 파라메타를 모두 불러오고, pretrained=False를 하면 ResNet18 구조만 불러옵니다. 이때 주의해야할 점은 모델과 텐서에 .to(device)를 붙여야만 GPU 연산이 가능하니 꼭 기입해야합니다.

model = torchvision.models.resnet18(pretrained=True)
print(model)

모델의 구조를 보면 마지막 출력 노드가 1000개라는 것을 알 수 있는데요, 이는 1000개의 클래스를 가진 ImageNet 데이터를 이용하여 사전학습 된 모델이기 때문입니다. 따라서 우리가 사용하는 CIFAR10 데이터에 맞게 출력층의 노드를 10개로 변경해야만 합니다.

기존의 모델

num_ftrs = model.fc.in_features # fc의 입력 노드 수를 산출한다. 512개
model.fc = nn.Linear(num_ftrs, 10) # fc를 nn.Linear(num_ftrs, 10)로 대체한다.
model = model.to(device)
print(model)

fc에서 out_feature가 10으로 변한것을 알 수 있다.

 

이어서 loss function과 optimizer를 정의하여 학습을 수행합니다. 코드는 다음과 같습니다.

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-2)
for epoch in range(20):

    running_loss = 0.0
    for data in trainloader:
        
        inputs, labels = data[0].to(device), data[1].to(device)
          
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    cost = running_loss / len(trainloader)        
    print('[%d] loss: %.3f' %(epoch + 1, cost))  

torch.save(model.state_dict(), './models/cifar10_resnet18.pth')      

print('Finished Training')
model = torchvision.models.resnet18(pretrained=False) # 모델의 구조를 먼저 불러오고 pre-trained된 파라메터를 덮어씌운다.
num_ftrs = model.fc.in_features # fc의 입력 노드 수를 산출한다. 512개
model.fc = nn.Linear(num_ftrs, 10) # fc를 nn.Linear(num_ftrs, 10)로 대체한다.
model = model.to(device)
model.load_state_dict(torch.load('./models/cifar10_resnet18.pth'))
correct = 0
total = 0
with torch.no_grad():
    model.eval()
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

---

모델 동결(Model Freezing)

 

피쳐 추출에 해당하는 합성곱 층의 변수를 업데이트 하지 않고 분류 파트에 해당하는 fully connected layer의 변수만 업데이트 할 수 있는데 이 때 변수가 업데이트 되지 않게 변수를 얼린다고 하여 이를 프리징(Freezing)이라고 합니다. 먼저 AlexNet 구조와 사전 학습 된 파라메타를 모두 불러오겠습니다. 모델의 구조를 보면 마지막 출력 노드가 1000개라는 것을 알 수 있는데요, 이는 1000개의 클래스를 가진 ImageNet 데이터를 이용하여 사전학습 된 모델이기 때문입니다. 따라서 이전 예제와 같이 우리가 사용하는 CIFAR10 데이터에 맞게 출력층의 노드를 10개로 변경해야만 합니다. 다음 모델은 AlexNet과 동일한 타입의 이미지 분류 문제를 위해 사용하였으므로, feature extraction층에 존재하는 가중치 값들은 동결하고 classifier층에 존재하는 가중치만을 재조정하여 customize된 작업에 사용할 수 있다.

model = torchvision.models.alexnet(pretrained=True)

num_ftrs = model.classifier[6].in_features # fc의 입력 노드 수를 산출한다. 
model.classifier[6] = nn.Linear(num_ftrs, 10) # fc를 nn.Linear(num_ftrs, 10)로 대체한다.
model = model.to(device)
# 합성곱 층은 0~9까지이다. 따라서 9번째 변수까지 역추적을 비활성화 한 후 for문을 종료한다.

for i, (name, param) in enumerate(model.named_parameters()):
    
    param.requires_grad = False
    if i == 9:
        print('end')
        break
# requires_grad 확인
print(model.features[0].weight.requires_grad) # False
print(model.features[0].bias.requires_grad) # False
print(model.features[3].weight.requires_grad) # False
print(model.features[3].bias.requires_grad) # False
print(model.features[6].weight.requires_grad) # False
print(model.features[6].bias.requires_grad) # False
print(model.features[8].weight.requires_grad) # False
print(model.features[8].bias.requires_grad) # False
print(model.features[10].weight.requires_grad) # False
print(model.features[10].bias.requires_grad) # False
print(model.classifier[1].weight.requires_grad) # True
print(model.classifier[1].bias.requires_grad) # True
print(model.classifier[4].weight.requires_grad) # True
print(model.classifier[4].bias.requires_grad) # True
print(model.classifier[6].weight.requires_grad) # True
print(model.classifier[6].bias.requires_grad) # True
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-2)
for epoch in range(20):

    running_loss = 0.0
    for data in trainloader:
        
        inputs, labels = data[0].to(device), data[1].to(device)
          
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    cost = running_loss / len(trainloader)        
    print('[%d] loss: %.3f' %(epoch + 1, cost))  
   

print('Finished Training')