[파이토치/Pytorch] FashionMNIST dataset으로 transfer learning 연습 + ray module로 hyperparameter tunning(예제 코드)

torchvision에서 제공하는 데이터셋 중 하나인 Fashion-Mnist Dataset으로 transfer-learning을 연습해본다. pre-trained 모델로는 imagenet_resnet18을 활용한다.

$ pip uninstall -y -q pyarrow
$ pip install -q -U ray[tune]
$ pip install -q ray[debug]

ray 활용을 위한 설치가 필요하다. debug 모드는 따로 없는 것 같다.

import torchvision
import torch

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
from tqdm.notebook import tqdm

# 데이터 로드
# ref - https://github.com/zalandoresearch/fashion-mnist
fashion_train = torchvision.datasets.FashionMNIST(root='./fashion', train=True, download=True)
fashion_test = torchvision.datasets.FashionMNIST(root='./fashion', train=False, download=True)

Fashion-Mnist의 이미지는 channel이 1개인 gray scale이다.

# ImageNet에서 학습된 ResNet 18 딥러닝 모델을 불러옴
imagenet_resnet18 = torchvision.models.resnet18(pretrained=True)
print("네트워크 필요 입력 채널 개수", imagenet_resnet18.conv1.weight.shape[1])
print("네트워크 출력 채널 개수 (예측 class type 개수)", imagenet_resnet18.fc.weight.shape[0])
print(imagenet_resnet18)

반면 활용하려는 모델의 경우 입력 채널은 3개가 필요하고, out_features는 1000개나 된다.

target_model = imagenet_resnet18
FASHION_INPUT_NUM = 1
FASHION_CLASS_NUM = 10

# target model의 입력 크기와 출력 크기를 변경하여 준다. 기존 부분 중 일부를 변경한다.
target_model.conv1 = torch.nn.Conv2d(FASHION_INPUT_NUM, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
target_model.fc = torch.nn.Linear(in_features=512, out_features=FASHION_CLASS_NUM, bias=True)

# 새롭게 넣은 네트워크 가중치를 xavier uniform으로 초기화
# (참고)왜 xavier uniform으로 초기화하는 이유 - 관련 논문을(https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)

# 앞에서 initalize 했던 distribution 외에도 자유롭게 채워 넣어봅시다.
torch.nn.init.xavier_uniform_(target_model.fc.weight)
# fully connected layer의 bias를 resnet18.fc in_feature의 크기의 1/root(n) 크기의 uniform 분산 값 중 하나로 설정하는 것이 좋다.
# (참고) https://stackoverflow.com/questions/49433936/how-to-initialize-weights-in-pytorch
stdv = 1/np.sqrt(512) 
target_model.fc.bias.data.uniform_(-stdv, stdv)


print("네트워크 필요 입력 채널 개수", target_model.conv1.weight.shape[1])
print("네트워크 출력 채널 개수 (예측 class type 개수)", target_model.fc.weight.shape[0])

이를 해결하기 위해 맨 앞의 layer와 fc layer에 수정을 해준다. 마지막 layer의 출력 노드의 수는 10개로, 첫 conv1 layer의 경우 입력 channel을 1개만 활용하도록 한다. 이때 w와 b에 초기화가 필요한데, 위 주석의 방법을 참고하자.

 

 

common_transform = torchvision.transforms.Compose(
  [
    torchvision.transforms.ToTensor() # PIL Image를 Tensor type로 변경함
  ]
)
# transform에는 다양한 전처리 혹은 Augmentation이 포함될 수 있다.
fashion_train_transformed = torchvision.datasets.FashionMNIST(root='./fashion', train=True, download=True, transform=common_transform)
fashion_test_transformed = torchvision.datasets.FashionMNIST(root='./fashion', train=False, download=True, transform=common_transform)

# DataLoader에 할당, gpu 사용시에는 num_workers를 4로 하자.
BATCH_SIZE = 64
fashion_train_dataloader = torch.utils.data.DataLoader(fashion_train_transformed, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
fashion_test_dataloader = torch.utils.data.DataLoader(fashion_test_transformed, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)


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

target_model.to(device) 

LEARNING_RATE = 0.0001 
NUM_EPOCH = 5 

# 분류 학습 때 많이 사용되는 Cross entropy loss를 objective function으로 사용
# (참고) https://en.wikipedia.org/wiki/Cross_entropy
loss_fn = torch.nn.CrossEntropyLoss() 
optimizer = torch.optim.Adam(target_model.parameters(), lr=LEARNING_RATE) 

dataloaders = {
    "train" : fashion_train_dataloader,
    "test" : fashion_test_dataloader
}

원하는 전처리를 적용한 dataloader를 생성하고, gpu 사용 여부, lr, epoch, optimizer 등 hyperparameter를 설정한다. 위처럼 설정하고 진행하면, tuning을 따로 하지 않는 경우이다. 위의 경우 말고, lr, epoch, batch_size에 대해 적절한 hyperparameter를 탐색하는 코드는 아래와 같다.

## 1. Learning Rate
def get_adam_by_learningrate(model, learning_rate:float):
  return torch.optim.Adam(model.parameters(), lr=learning_rate)
## 2. Epoch 개수
def get_epoch_by_epoch(epoch:int):
  return epoch
## 3. BatchSize 크기에 따른 데이터 로더 생성
common_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
fashion_train_transformed = torchvision.datasets.FashionMNIST(root='./fashion', train=True, download=True, transform=common_transform)
fashion_test_transformed = torchvision.datasets.FashionMNIST(root='./fashion', train=False, download=True, transform=common_transform)

def get_dataloaders_by_batchsize(batch_size:int):
  # Mnist Dataset을 DataLoader에 붙이기
  BATCH_SIZE = batch_size
  fashion_train_dataloader = torch.utils.data.DataLoader(fashion_train_transformed, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
  fashion_test_dataloader = torch.utils.data.DataLoader(fashion_test_transformed, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

  dataloaders = {
      "train" : fashion_train_dataloader,
      "test" : fashion_test_dataloader
  }

  return dataloaders

from ray import tune

config_space = {
    "NUM_EPOCH" : tune.choice([4,5,6,7,8,9]),
    "LearningRate" : tune.uniform(0.0001, 0.001),
    "BatchSize" : tune.choice([32,64,128]),
}
loss_fn = torch.nn.CrossEntropyLoss()

다음으로 ray가 탐색할 파라미터와 그 범위를 지정한다. tune.choice를 통해 1개를 고르거나, tune.uniform을 통해 해당 범위에서 하나를 랜덤하게 고를 수 있다.

 

def training(
    config # 조작 변인 learning rate, epoch, batchsize 정보
):
  # 통제 변인
  target_model = get_imagenet_pretrained_model() 

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

  # 조작 변인
  NUM_EPOCH = get_epoch_by_epoch(config["NUM_EPOCH"])
  dataloaders = get_dataloaders_by_batchsize(config["BatchSize"])
  optimizer = get_adam_by_learningrate(target_model, config["LearningRate"])

  ### 학습 코드 시작
  best_test_accuracy = 0.
  best_test_loss = 9999.
  for epoch in range(NUM_EPOCH):
    for phase in ["train", "test"]:
      running_loss = 0.
      running_acc = 0.
      if phase == "train":
        target_model.train() 
      elif phase == "test":
        target_model.eval() 

      for ind, (images, labels) in enumerate(tqdm(dataloaders[phase])):
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad() # parameter gradient를 업데이트 전 초기화함

        # train 모드일 시에는 gradient를 계산하고, 아닐 때는 gradient를 계산하지 않아 연산량 최소화
        with torch.set_grad_enabled(phase == "train"): 
          logits = target_model(images)
          # 모델에서 linear 값으로 나오는 예측 값 ([0.9,1.2, 3.2,0.1,-0.1,...])을 최대 output index를 찾아 예측 레이블([2])로 변경함
          _, preds = torch.max(logits, 1)   
          loss = loss_fn(logits, labels)

          if phase == "train":
            loss.backward() # 모델의 예측 값과 실제 값의 CrossEntropy 차이를 통해 gradient 계산
            optimizer.step() 

        running_loss += loss.item() * images.size(0) # 한 Batch에서의 loss 값 저장
        running_acc += torch.sum(preds == labels.data) # 한 Batch에서의 Accuracy 값 저장

      # 한 epoch이 모두 종료되었을 때,
      epoch_loss = running_loss / len(dataloaders[phase].dataset)
      epoch_acc = running_acc / len(dataloaders[phase].dataset)

      if phase == "test" and best_test_accuracy < epoch_acc: 
        best_test_accuracy = epoch_acc
      if phase == "test" and best_test_loss > epoch_loss: #
        best_test_loss = epoch_loss
  # epoch 종료
  tune.report(accuracy=best_test_accuracy.item(), loss=best_test_loss)

나머지는 일반적인 학습 과정과 거의 동일하지만, tune.report를 사용해 이후 tune.run을 통해 ray가 성능을 비교할 수 있도록 한다.

from ray.tune.suggest.hyperopt import HyperOptSearch

# HyperOptSearch 통해 Search를 진행
# 더 다양한 Optimizer들은 https://docs.ray.io/en/master/tune/api_docs/suggestion.html#bayesopt 참고
optim = HyperOptSearch( 
    metric='accuracy', # hyper parameter tuning 시 최적화할 metric을 결정
    mode="max", # target objective를 maximize 하는 것을 목표로
)

어떤 metric을 가지고 어떤 목적을 수행할 지를 설정한 뒤 아래 코드를 실행하면 된다.

from ray.tune import CLIReporter
import ray

NUM_TRIAL = 3 # Hyper Parameter를 탐색할 때에, 실험을 최대 수행할 횟수를 지정

# CLIReporter는 중간 수행 결과를 command line에 출력하도록 함
reporter = CLIReporter( 
    parameter_columns=["NUM_EPOCH", "LearningRate", "BatchSize"],
    metric_columns=["accuracy", "loss"])

ray.shutdown() # ray 초기화 후 실행

analysis = tune.run(
    training,
    config=config_space,
    search_alg=optim,
    #verbose=1,
    progress_reporter=reporter,
    num_samples=NUM_TRIAL
    # Colab 런타임이 GPU를 사용하지 않는다면 아래는 주석 처리
    #resources_per_trial={'gpu': 1}
)

tune.run이 파라미터 탐색을 시작하는 코드다.

이런 느낌의 중간 과정들을 보여주다가,

탐색을 마친다.

best_trial = analysis.get_best_trial('accuracy', 'max')
print(f"최고 성능 config : {best_trial.config}")
print(f"최고 test accuracy : {best_trial.last_result['accuracy']}")

위 코드로 최고 성능의 모델을 찾아보면,

결과를 출력해준다.