WandB(Weights & Biases)로 하이퍼 파라미터 튜닝하기

모델의 성능을 개선하기 위해 다양한 하이퍼 파라미터를 조정하고, 여러 번의 실험을 반복하며 모델의 동작을 이해하는 것은 필수적이다. 이를 위해 Weights & Biases(WandB)라는 도구를 사용하면 효과적으로 모델 실험을 관리하고 하이퍼 파라미터 튜닝을 할 수 있다. 이번 글에서 WandB에 대한 간략한 소개와, PyTorch 프레임워크에서 WandB Sweep을 사용하는 방법에 대해 알아보자.

---

WandB(Weights & Biases)란?

WandB(Weights & Biases)는 머신러닝 실험 관리를 위한 오픈소스 도구로, 실험의 결과와 메트릭을 추적하고 시각화하는 기능을 제공한다. 또한, 하이퍼파라미터 튜닝을 위한 Sweep 기능을 제공하여 자동화된 하이퍼파라미터 검색을 수행할 수 있다. WandB를 사용하면 실험의 성과를 효과적으로 모니터링하고, 실험 결과를 비교하고, 모델 개선을 위한 최적의 하이퍼 파라미터를 찾는 등의 작업을 간편하게 수행할 수 있다.

WandB Sweep이란?

WandB Sweep는 하이퍼 파라미터 튜닝을 위한 자동화된 기능으로, 여러 가지 하이퍼파라미터 조합을 시도하여 최적의 조합을 찾는 기능이다. 사용자가 하이퍼 파라미터의 값 범위를 지정하고, 실험을 수행할 수 있는 소스코드를 작성하면, WandB가 자동으로 다양한 조합을 시도하고, 실험 결과를 기록하여 최적의 하이퍼 파라미터 조합을 찾을 수 있다.

PyTorch에서 WandB Sweep 사용하기

PyTorch에서 WandB Sweep을 사용하는 방법을 간단한 예제를 통해 알아보자.

1. WandB 설치하기

먼저, WandB를 설치해야 한다. 아래의 명령어를 사용하여 WandB를 설치한다.

pip install wandb

 

2. WandB 프로젝트 초기화하기

WandB 프로젝트를 초기화하려면, WandB 웹사이트(https://wandb.ai/)에서 계정을 만들고 로그인해야 한다. 그리고 프로젝트 디렉토리에서 아래의 명령어를 실행하여 WandB 프로젝트를 초기화한다.

wandb init

위와 같이 API key를 통해 인증을 요청한다. 메세지의 내용에 따라 https://wandb.ai/authorize로 접속해 API key를 복사해서 붙여 넣어주면 WandB 프로젝트를 초기화할 수 있다.

 

3. 실험 소스코드 작성하기

하이퍼 파라미터 튜닝을 위한 실험 소스코드를 작성한다. 이 소스코드에는 실험할 모델, 데이터 로딩, 모델 학습 및 평가 등의 과정을 포함해야 한다. WandB에서 제공하는 wandb.init() 함수를 사용하여 실험을 초기화하고, 실험 설정과 하이퍼파라미터 값을 저장할 수 있다.

import wandb
import torch
from transformers import AutoModelForSequenceClassification,TrainingArguments, Trainer

def model_init():
    model = AutoModelForSequenceClassification.from_pretrained("monologg/koelectra-base-finetuned-nsmc",num_labels=1,ignore_mismatched_sizes=True)
    return model

def train(config=None):
    with wandb.init(config=config): # WandB 초기화
        # 하이퍼 파라미터 가져오기
        config = wandb.config
        t = config.learning_rate

        # train에 필요한 파라미터 지정
        args = TrainingArguments(
            output_dir=f"./checkpoint/baseline_Test_fine_{t}",
            evaluation_strategy="epoch",
            save_strategy="epoch",
            report_to='wandb',
            learning_rate=config.learning_rate,
            per_device_train_batch_size=config.batch_size,
            per_device_eval_batch_size=config.batch_size,
            num_train_epochs=config.epochs,
            weight_decay=config.weight_decay,
            load_best_model_at_end=True,
            dataloader_num_workers=4,
            logging_steps=200,
        )

        # hugging face Trainer를 활용하여 모델 학습 및 평가
        trainer = Trainer(
            model_init = model_init,
            args=args,
            train_dataset=Train_textDataset,
            eval_dataset=Val_textDataset,
            compute_metrics=compute_pearson_correlation
        )

        trainer.train()

 

4. WandB Sweep 작성

하이퍼 파라미터 튜닝을 위한 Sweep 설정을 작성한다. Sweep 설정은 YAML 파일로 작성하거나, 코드 내에서 딕셔너리 형태로 지정할 수 있다. Python 코드 내에서 wandb.sweep() 함수를 사용하여 하이퍼 파라미터 조건을 지정하고, sweep_id를 얻을 수 있다. 아래는 예제 코드이다.

parameters_dict = {
    'epochs': {
        'value': 8
    },
    'batch_size': {
        'values': [4,8,16]
    },
    'learning_rate': {
        'distribution': 'log_uniform_values',
        'min': 1e-5,
        'max': 5e-5
    },
    'weight_decay': {
        'values': [0.3,0.4,0.5]
    },
}

# 하이퍼 파라미터 조건 지정
sweep_config = {
    'method': 'random',
    'parameters': parameters_dict,
}

# Sweep을 설정하는 코드
sweep_id = wandb.sweep(sweep=sweep_config, project='my_project_name')

이 설정에서는 epoch를 8로 고정하고, lr 하이퍼파라미터의 값을 0.00001에서 0.00005 사이에서, batch_size 하이퍼파라미터의 값을 4, 8, 16 중에서, weight_decay를 0.3, 0.4, 0.5 중에서 랜덤하게 선택하여 실험을 수행한다.

 

5. Sweep을 실행하는 코드 작성

Sweep을 실행하는 코드를 작성한다. sweep_id를 사용하여 Sweep을 시작하고, 사전에 정의한 하이퍼 파라미터 범위에서 조합을 만들어 실험을 수행한다.

# Sweep을 실행하는 코드
wandb.agent(sweep_id, function=train, count=10)

위 코드에서 function은 실행할 함수를 지정하고, count는 실험을 실행할 횟수를 지정한다.

 

6. WandB Sweep 결과 확인하기

실행된 WandB Sweep의 결과는 WandB 웹 대시보드에서 확인할 수 있다. 각 실험의 결과 및 메트릭들이 시각화되어 편리하게 확인할 수 있다.

 

• Parallel coordinates

다양한 하이퍼 파라미터들과 모델 메트릭을 한 눈에 살펴볼 수 있다.

• scatter plot

Sweep 동안 생성된 여러 Run들을 비교하여 실험이 진행되는 상황을 시각화할 수 있다.

• parameter importance

지금까지 진행된 실험을 바탕으로 메트릭 값과 높은 연관성이 있는 하이퍼 파라미터를 보여준다.

 

7. 최적의 하이퍼파라미터 조합으로 실험 재현하기

WandB Sweep을 통해 찾은 최적의 하이퍼파라미터 조합으로 실험을 재현하기 위해서는 WandB 웹 대시보드에서 해당 실험의 하이퍼파라미터 값을 확인하고, 소스코드 내에서 해당 값으로 설정한 뒤 실험을 다시 실행하여 최적의 모델을 얻을 수 있다.

프로젝트 대시보드에서 좌측의 Runs 중 원하는 Run을 누른 뒤,

 

좌측 메뉴바 가장 위에 위치한 Overview를 누른 다음

 

페이지 하단 config에서 원하는 파라미터를 검색하여 값을 확인할 수 있다.

 

+) 추가

• sweep configuration의 구조

https://docs.wandb.ai/guides/sweeps/define-sweep-configuration

Define sweep configuration | Weights & Biases Documentation