Timm을 이용한 Finetuning 실습

이번에는 이전에 배웠던 Timm 라이브러리의 사전 학습 모델을 이용해 보는 실습을 진행하였다. 

1. 라이브러리 import 

import torch 
import numpy as np
import warnings
import matplotlib.pyplot as plt 
import pandas as pd 
from sklearn.model_selection import train_test_split 
from sklearn.metrics import accuracy_score
from tqdm.notebook import tqdm 

warnings.filterwarnings('ignore')

import torch.nn as nn # 모델 구성
from torchvision.datasets import CIFAR10 # CIFAR10 데이터셋 불러오는 라이브러리
import torchvision.transforms as T # 이미지 변환
import torch.optim as optim # optimizer 설정

import random
import torch.backends.cudnn as cudnn

def random_seed(seed_num):
    torch.manual_seed(seed_num)
    torch.cuda.manual_seed(seed_num)
    torch.cuda.manual_seed_all(seed_num)
    np.random.seed(seed_num)
    cudnn.benchmark = False
    cudnn.deterministic = True
    random.seed(seed_num)
random_seed(624)

device = 'cuda:0'

위와 같이 실습에 필요한 라이브러리들을 불러와준다. 현재 `python=3.9.0`, `Pytorch=2.0.0`, `torchvision=0.15.0`에서 작업하고 있으며 GPU를 활용할 것이다. 

2. 모델 불러오기

다음으로 Timm 라이브러리를 설치하고 import하고, 모델 리스트를 확인해 볼 수 있다. 

import timm
# 전체 모델 검색
timm.list_models() 
# 쿼리를 통해 특정 모델 검색 
timm.list_models('resnet*')
# pretrained weight가 존재하는 모델 검색
timm.list_models('resnet50', pretrained=True)

이때 위의 models들은 모두 pretrained 가중치를 가지고 있는 것이 아니기 때문에 전이학습을 위해 해당 모델을 찾기 위해서는 `pretrained=True`라는 옵션을 추가하여 검색해줘야 한다. 

# 모델 불러오기
model = timm.create_model('resnet50', pretrained=True)
# 모델 정보 확인
model.default_cfg
# 모델 구조 확인
model

위의 모델을 이용하여 CIFAR10 데이터셋을 평가해볼 것인데, 위의 우측 사진을 보면 기존 모델 구조의 마지막 fully connnectd layer을 보면 1000개를 구분하는 것을 볼 수 있다. 이는 이 모델이 ImageNet 데이터셋을 학습한 데이터 셋이기 때문인데 CIFAR10은 10가지로 분류해야 하기 때문에 이 부분을 수정해야 한다. 이것도 Timm에서 간단하게 제공하는데, 모델을 불러올 때 `num_classes=10`라는 옵션을 넣어주기만 하면 된다.(단 이렇게 불러온 모델의 마지막 layer은 학습되지 않은 랜덤 한 가중치 상태이기 때문에 학습이 필요하다.)

model2 = timm.create_model('resnet50', pretrained = True, num_classes = 10)

3. 데이터 불러오기

다음으로 CIFAR10 데이터셋을 불러올 것이다. 데이터셋에 대하여 간단하게 설명하자면 총 10개의 클래스( Airplane , automobile, bird, cat, deer, dog, frog, horse, ship,truck)을 가지며, 5만 개의 학습 데이터와 1만 개의 테스트 데이터로 구성되어 있다. 

# 데이터 불러오기
cifar_transform = T.Compose([
    T.ToTensor(), # 텐서 형식으로 변환
])
download_root = './CIFAR10_DATASET'

trainval_dataset = CIFAR10(download_root, transform=cifar_transform, train=True, download=True) # train dataset 다운로드
test_dataset = CIFAR10(download_root, transform=cifar_transform, train=False, download=True) # test dataset 다운로드

train_num, valid_num = int(len(trainval_dataset) * 0.8), int(len(trainval_dataset) * 0.2) # 8 : 2 = train : valid
print("Train dataset 개수 : ",train_num)
print("Validation dataset 개수 : ",valid_num)
train_dataset,val_dataset = torch.utils.data.random_split(trainval_dataset, [train_num, valid_num]) # train - valid set 나누기

데이터를 tensor형태로 불러오고 train 데이터셋에서 valid 셋을 나누어준다. 

다음은 DataLoader을 이용하여 배치단위의 데이터셋을 불러올 것이다. 

BATCH_SIZE = 64 
train_dataloader = torch.utils.data.DataLoader(dataset=train_dataset,
                                          batch_size=BATCH_SIZE,
                                          shuffle=True,
                                          drop_last=False, num_workers = 8) # train dataloader 구성
val_dataloader = torch.utils.data.DataLoader(dataset=val_dataset,
                                          batch_size=BATCH_SIZE,
                                          shuffle=False,
                                          drop_last=False, num_workers = 8) # valid dataloader 구성
test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=BATCH_SIZE,
                                          shuffle=False,
                                          drop_last=False, num_workers = 8) # test dataloader 구성

4. 중간 추론

이제 모델을 이용하여 하나의 데이터를 임의로 추론해보자. 

우선 GPU를 이용하기 위해 모델을 GPU로 올려주고 추가학습이 진행되지 않은 상태로 하나의 데이터셋을 추론하였다. 

device = 'cuda:0' 
model = timm.create_model('resnet50', pretrained=True, num_classes = 10).to(device) 

img, label = train_dataset[0]
img = img.unsqueeze(0) # 배치 추가

model.eval() # evaluation 상태로 만듦 (freeze)
preds = model(img.to(device)) # model inference, image 도 gpu에 올리기
pred_label = torch.argmax(preds).item() # 가장 큰 값의 index 반환
print(f'True Label : {label} \nPredict Label : {pred_label}')

# >> True Label:6
# >> Pred Label:7

위와 같이 pretrained model로도 예측이 잘못된 이유는 이전에 fully connectedlayer을 기존의 1000개가 아닌 10개로 바꾸면서 랜덤 한 가중치로 초기화되었기 때문에 추가학습이 반드시 필요한 상태이다. 

5. Fine Tuning 

PyTorch를 이용하여 모델을 학습하는 내용의 코드는 이전의 https://changsroad.tistory.com/457 에서 설명했기 때문에 아래를 펼쳐보면 training, evaluation, epoch 반복에 대한 코드가 작성되어 있다. 

# training 코드, evaluation 코드, training_loop 코드
def training(model, dataloader, train_dataset, criterion, optimizer, device, epoch, num_epochs):
  model.train()  # 모델을 학습 모드로 설정
  train_loss = 0.0
  train_accuracy = 0

  tbar = tqdm(dataloader)
  for images, labels in tbar:
      images = images.to(device)
      labels = labels.to(device)

      # 순전파
      outputs = model(images)
      loss = criterion(outputs, labels)

      # 역전파 및 가중치 업데이트
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      # 손실과 정확도 계산
      train_loss += loss.item()
      # torch.max에서 dim 인자에 값을 추가할 경우, 해당 dimension에서 최댓값과 최댓값에 해당하는 인덱스를 반환
      _, predicted = torch.max(outputs, 1)
      train_accuracy += (predicted == labels).sum().item()

      # tqdm의 진행바에 표시될 설명 텍스트를 설정
      tbar.set_description(f"Epoch [{epoch+1}/{num_epochs}], Train Loss: {loss.item():.4f}")

  # 에폭별 학습 결과 출력
  train_loss = train_loss / len(dataloader)
  train_accuracy = train_accuracy / len(train_dataset)

  return model, train_loss, train_accuracy

def evaluation(model, dataloader, val_dataset, criterion, device, epoch, num_epochs):
  model.eval()  # 모델을 평가 모드로 설정
  valid_loss = 0.0
  valid_accuracy = 0

  with torch.no_grad(): # model의 업데이트 막기
      tbar = tqdm(dataloader)
      for images, labels in tbar:
          images = images.to(device)
          labels = labels.to(device)

          # 순전파
          outputs = model(images)
          loss = criterion(outputs, labels)

          # 손실과 정확도 계산
          valid_loss += loss.item()
          # torch.max에서 dim 인자에 값을 추가할 경우, 해당 dimension에서 최댓값과 최댓값에 해당하는 인덱스를 반환
          _, predicted = torch.max(outputs, 1)
          valid_accuracy += (predicted == labels).sum().item()

          # tqdm의 진행바에 표시될 설명 텍스트를 설정
          tbar.set_description(f"Epoch [{epoch+1}/{num_epochs}], Valid Loss: {loss.item():.4f}")

  valid_loss = valid_loss / len(dataloader)
  valid_accuracy = valid_accuracy / len(val_dataset)

  return model, valid_loss, valid_accuracy


def training_loop(model, train_dataloader, valid_dataloader, train_dataset, val_dataset, criterion, optimizer, device, num_epochs, patience, model_name):
    best_valid_loss = float('inf')  # 가장 좋은 validation loss를 저장
    early_stop_counter = 0  # 카운터
    valid_max_accuracy = -1

    for epoch in range(num_epochs):
        model, train_loss, train_accuracy = training(model, train_dataloader, train_dataset, criterion, optimizer, device, epoch, num_epochs)
        model, valid_loss, valid_accuracy = evaluation(model, valid_dataloader, val_dataset, criterion, device, epoch, num_epochs)

        if valid_accuracy > valid_max_accuracy:
          valid_max_accuracy = valid_accuracy

        # validation loss가 감소하면 모델 저장 및 카운터 리셋
        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            torch.save(model.state_dict(), f"./model_{model_name}.pt")
            early_stop_counter = 0

        # validation loss가 증가하거나 같으면 카운터 증가
        else:
            early_stop_counter += 1

        print(f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f} Valid Loss: {valid_loss:.4f}, Valid Accuracy: {valid_accuracy:.4f}")

        # 조기 종료 카운터가 설정한 patience를 초과하면 학습 종료
        if early_stop_counter >= patience:
            print("Early stopping")
            break

    return model, valid_max_accuracy

5-1. 전체 학습

우선 첫 번째로 마지막 fully connected layer만 학습시키는 것이 아닌 모델 전체의 가중치를 학습시키는 방법을 이용해 보았다. 

# 모델 전체 fine tuning
num_epochs = 100
patience = 3
scores = dict()
model_name = 'exp1'

lr = 1e-3
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr = lr)
model, valid_max_accuracy = training_loop(model, train_dataloader, val_dataloader, train_dataset, val_dataset, criterion, optimizer, device, num_epochs, patience, model_name)
scores[model_name] = valid_max_accuracy

model.load_state_dict(torch.load("./model_exp1.pt")) # 모델 불러오기
model = model.to(device)
model.eval()
total_labels = []
total_preds = []
with torch.no_grad():
    for images, labels in tqdm(test_dataloader):
        images = images.to(device)
        labels = labels

        outputs = model(images)
        # torch.max에서 dim 인자에 값을 추가할 경우, 해당 dimension에서 최댓값과 최댓값에 해당하는 인덱스를 반환
        _, predicted = torch.max(outputs.data, 1)

        total_preds.extend(predicted.detach().cpu().tolist())
        total_labels.extend(labels.tolist())

total_preds = np.array(total_preds)
total_labels = np.array(total_labels)
full_model_tuning_acc = accuracy_score(total_labels, total_preds) 
print("Full Fine tuning model accuracy : ",full_model_tuning_acc) 
# >> Full Fine tuning model accuracy :  0.8239

결과적으로 전체 모델을 학습하여 진행하였더니 accuracy가 0.8239라는 수치가 나왔다. 

5-2. 일부 학습

다음은 마지막 fully connected layer만 학습시키는 방법을 사용해볼 것이다.

# 마지막 layer 만 fine tuning
num_epochs = 100
patience = 3
scores = dict()
model_name = 'exp2'

model = timm.create_model('resnet50', pretrained=True, num_classes= 10).to(device)

for para in model.parameters(): # 모든 layer freeze 하기
    para.requires_grad = False
for para in model.fc.parameters(): # fc layer 만 학습하기
    para.requires_grad = True

lr = 1e-3
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr = lr)
model, valid_max_accuracy = training_loop(model, train_dataloader, val_dataloader, train_dataset, val_dataset, criterion, optimizer, device, num_epochs, patience, model_name)
scores[model_name] = valid_max_accuracy

코드에서 볼 수 있듯이 우선 사전학습된 모델을 불러오고 fc layer을 제외한 다른 layer의 가중치는 모두 freeze 해준다. 이후 동일하게 학습을 진행한다. 또한 동일한 방법으로 추론까지 진행해 주었다.

model.load_state_dict(torch.load("./model_exp2.pt")) # 모델 불러오기
model = model.to(device)
model.eval()
total_labels = []
total_preds = []
with torch.no_grad():
    for images, labels in tqdm(test_dataloader):
        images = images.to(device)
        labels = labels

        outputs = model(images)
        # torch.max에서 dim 인자에 값을 추가할 경우, 해당 dimension에서 최댓값과 최댓값에 해당하는 인덱스를 반환
        _, predicted = torch.max(outputs.data, 1)

        total_preds.extend(predicted.detach().cpu().tolist())
        total_labels.extend(labels.tolist())

total_preds = np.array(total_preds)
total_labels = np.array(total_labels)
fc_tuning_acc = accuracy_score(total_labels, total_preds) 
print("Only FC Layer Fine tuning model accuracy : ",fc_tuning_acc)
# >> Only FC Layer Fine tuning model accuracy :  0.444

그 결과 accuracy가 0.444가 나왔다.

6. 결과 비교

우선 사전학습 모델을 불러온 다음 fully connected layer을 포함한 전체 weight를 재학습시켰을 때는 accuracy가 `0.8239`가 나왔으며, 다른 layer들의 가중치는 고정시키고 fully connected layer만 학습시켰을 때는 `0.444`로 더 안 좋은 정확도가 계산됐다. 이 이유는 데이터셋의 크기가 40000개 정도로 매우 크기 때문에 전체 재학습 혹은 layer을 추가하여 재학습시키는 방법이 더 좋은 것이다.(보통 새로운 데이터셋이 작은데 유사할 경우에 마지막 fc만 학습) 

이렇게 모델을 재학습시키는 방법을 이용해서 learning rate 등을 수정해 보면서 정확도를 개선시켜 볼 수 있다.