[WandB] Huggingface 라이브러리에서 Sweep 사용하기

서론 (생략 가능)

 WandB를 활용하면 실험 기록 등을 기록하는데 간편하고, 특히 팀프로젝트 등을 진행할 때 실험 관리를 하기 매우 수월하다. 또한, 한국어 Document도 지원하기 때문에 처음 사용하는 사람들도 쉽게 사용할 수 있다고 생각한다. 하지만 비교적 최근에 생긴 Tool이다 보니 공식 Document이외에 정보가 매우 한정적이었다.

 이번에 내가 겪었던 문제도 Huggingface에서 Sweep을 하고자 하니 정보가 많이 부족해서 생겼다. 우선, 공식 Document의 Sweep을 보면 Sweep에 대한 설명이 자세하게 잘 되어있다. 사용법과 간략한 예시 코드도 꽤 잘 정리되어 있다. 따라서 만약 Huggingface를 사용하지 않는 사용자라면 공식 Document만 봐도 잘 할 수 있을 것이라고 생각한다.  

 하지만 만약 우리처럼 Huggingface라이브러리로 구현된 경우 Sweep을 이용하기 까다로울 것이라고 생각된다. 실제 예시 코드를 보더라도, 내가 찾았던 정보 중에는, Huggingface로 구현된 예시를 찾을 수 없었기 때문이다. 따라서 꽤 고생을 했고, 내 방법이 정답은 아닐지라도 다른 분들에게 조금이라도 도움이 됐으면 한다.

Sweep 사용법 (Basic)

우선 Sweep을 사용하기 위해서는 사용방법을 알아야 한다. Sweep은 크게 dictionary 정의, Sweep 구성, Sweep을 초기화, 에이전트를 실행 과정으로 진행된다. (wandb login 과정은 생략하겠습니다_해당 포스팅 참조)

1. Python train 스크립트 설정

import wandb

# wandb.init 전에 default hyperparameter를 세팅합니다.
# 해당 parameter들이 sweep에 적절하게 설정됩니다. 
hyperparameter_defaults = dict(
    dropout = 0.5,
    channels_one = 16,
    channels_two = 32,
    batch_size = 100,
    learning_rate = 0.001,
    epochs = 2,
    )

# default값들을 wandb.init에 넣어줍니다.
wandb.init(config=hyperparameter_defaults)
config = wandb.config

# Your model here ... (모델을 정의하고)

# Log metrics inside your training loop (매트릭을 정의해서 log를 기록합니다.)
metrics = {'accuracy': accuracy, 'loss': loss}
wandb.log(metrics)

해당 단계는 기본적으로 sweep을 사용하기에 앞서 train.py를 설정하는 부분입니다. 사실 huggingface에서는TrainingArguments에 report_to="wandb"만 넣어주면 쉽게 wandb를 사용할 수 있습니다. 따라서 wandb.log를 사용하지도 않고, wandb.init도 크게 중요한 부분은 아닐 수 있습니다. 하지만 여기서 config=wandb.config를 설정하는 부분이중요합니다. 해당 부분을 설정해야만, sweep으로 새로 정의된 파라미터 값들이 해당 위치에 제대로 적용될 수 있기 때문입니다.

(추가적으로 제가 겪었던 오류를 말씀드리자면, config를 설정하지 않아 10번의 task를 진행하는 동안 똑같은 10개의 모델을 학습시킨 적이 있습니다. 따라서 저 부분에 존재를 아셨다면 이미 반은 앞서가셨다고 볼 수 있습니다)

 

2. Sweep Config (스윕 구성)

 이제 YAML파일을 설정해야 하는데, 공식 Document에는 다음과 같이 쓰여있다. 'YAML 파일을 설정하여 훈련 스크립트, 매개변수 범위(parameter ranges), 검색 전략(search strategy) 및 중지 기준(stopping criteria)을 지정합니다. W&B는 이러한 매개변수와 해당 값을 명령줄 전달인자(command line arguments)로 훈련 스크립트에 전달하며 config 객체를 사용해 자동으로 분석합니다.' 

 여기서 config 객체를 사용해 자동으로 분석한다는 부분이 핵심이라는 것을 기억하자. 아래는 sweep.yaml의 예시이다. 이렇게 작성된 parameter들이 앞서 정의한 wandb.config에 할당된다고 생각하면 된다. 

# sweep.yaml
program: train.py
method: bayes
metric:
  name: val_loss
  goal: minimize
parameters:
  learning_rate:
    min: 0.001
    max: 0.1
  optimizer:
    values: ["adam", "sgd"]

만약, 최적화할 메트릭을 지정하는 경우 메트릭을 로그했는지 확인해야 한다. 위 예시에서, 구성(config) 파일에 val_loss가 있으므로, 스크립트에 정확한 메트릭 이름을 로그해야 한다고 공식 Docu에서 얘기한다. [wandb.log({"val_loss": validation_loss})] 하지만, Huggingface에서 metric을 사용할 때는 'eval/F1'처럼 사용한다.

3. Sweep 초기화

 앞서 정의한 스윕 구성 파일(sweep config file)을 설정하고, 다음의 명령을 실행하여 시작한다.

wandb sweep sweep.yaml

해당 코드를 터미널에서 실행하면 Entity(사용자이름), Project 이름을 포함한 Sweep ID를 출력한다. 이 SweepID를 복사해두자.

 

4. Agent실행

 위애서 Sweep을 초기화하고 나온 ID를 활용해서 wandb agent SweepID 형태로 입력해주면 Agent를 실행하게 된다. 

wandb agent your-sweep-id

Agent가 실행되게 되면 앞서 정의한 프로젝트에서 Sweep이 실행됨을 확인할 수 있을 것이다. 

Sweep 사용법 (Huggingface)

 Huggingface라이브러리에서 Sweep을 사용하기 위해 기본적으로 Sweep Config를 만들고, Sweep ID를 생성하고, Agent를 실행하는 과정까지는 비슷하다. 하지만, Huggingface 라이브러리에서 내가 실험하고자 하는 변수들을 어떻게 설정할 것인가 하는 문제점에 대해서는 해답을 얻기 어려웠을 것이라고 생각한다. 따라서 해당 부분을 코드와 함께 설명하도록 하겠다.

# train.py
from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification, Trainer, TrainingArguments
import wandb
from ray.tune.integration.wandb import wandb_mixin

def compute_metrics(pred):
  """ validation을 위한 metrics function """
  ...
  return {
      'micro f1 score': f1,
      'auprc' : auprc,
      'accuracy': acc,
  }
### 1 ###
@wandb_mixin 
train():
...
  wandb.init(project=project)
  config = wandb.config ### 2 ###
  training_args = TrainingArguments(
	output_dir='./results',          # output directory
	save_total_limit=5,              # number of total save model.
	save_steps=500,                 # model saving step.
    ### 3 ###
	num_train_epochs=config["num_train_epochs"],         # total number of training epochs
	learning_rate=config["learning_rate"],               # learning_rate
	per_device_train_batch_size=32,  # batch size per device during training
	per_device_eval_batch_size=16,   # batch size for evaluation
	...
	report_to="wandb",  # enable logging to W&B
                 	)
  trainer = Trainer(
    model=model,                     
    args=training_args,              
    train_dataset=train_dataset,      
    eval_dataset=dev_dataset,            
    compute_metrics=compute_metrics, ### 4 ###
    )
    trainer.train()

if __name__ == '__main__':
  train()

1. wandb_mixin데코레이터 설정

우선 @wandb_mixin데코레이터를 우리의 train() 함수 위에다가 정의해줘야 한다. Trainable 클래스 또는 함수 API에 대한 @wandb_mixin과 함께 사용할 Wandb API 초기화에 도움을 주는데 사용된다. train()함수 위에 다음과 같이 적어두기만 함녀 된다.

 

2. wandb.config 설정

 다음으로 wandb config값을 받아와야 한다. 이렇게 해야만 3.번에서 sweep.yaml 에서 정의한 파라미터 값들을 사용할 수 있다. 이걸 설정하지 않으면 계속 같은 모델을 학습하게 된다.

 

3. wandb.config값을 통해 파라미터 세팅

 하이퍼파라미터 튜닝을 하고 싶은 매개변수들에 대해 config["매개변수 값"]을 넣어준다.

 

4. sweep.yaml값 변경

 사실 1~3까지만 해도 wandb에서 sweep을 돌릴 수 있다. 추가적으로 sweep.yaml에서 metric을 설정하는 부분이 있는데, 이는 sweep이 달성해야하는 기준을 설정하는 부분이다. 위에서 설명했던 것처럼 기준으로 삼을 metric의 값이 변경됐다면 (예를들어 "micro f1 score") Huggingface에서는 eval/micro f1 score와 같이 metric값을 세팅해줘야 한다.

# sweep.yaml
program: train.py
method: bayes
metric:
  name: eval/micro f1 score
  goal: maximize
parameters:
  learning_rate:
    values: [1e-5, 3e-5, 5e-5]
  num_train_epochs:
    values: [3, 4, 5]
  per_device_train_batch_size:
    values: [32, 16]

 

 이렇게 하면 다음과 같이 WandB Document에서만 보던 그래프를 그릴 수가 있다...정작 해보면 별 거 아닌 것 같은데 관련된 설명이 없어서 꽤 고생을 많이 했었다. 물론 지금 내가 구현한 코드가 정답이 아닐 가능성도 매우 높다. 따라서 더 괜찮은 방법이나 잘못된 부분이 있다면 말씀해주시면 좋을 것 같다. 

 

 22-04-29 / 추가) Tuning Method가 'bayes'로 되어있는데, 이외에 'grid'나 'random'도 있다. 'bayes'검색방법은 가우스 프로세스를 사용해 매개변수와 모델 메트릭 사이의 관계를 모델링하고, 매개변수를 선택해 개선확률을 최적화하는 방법이라고 한다. 하지만, 해당 파라미터를 지정할 때, config 변수 값을 min, max 등 float값을 활용해야 하는 것으로 보인다. 그냥 values로 받게 될 경우 동일한 파라미터가 들어가게 되는 것처럼 보인다. 

참고) Sweep Configuration

 

 

참고 링크

• https://docs.wandb.ai/v/ko/sweeps
• https://docs.wandb.ai/v/ko/integrations/ray-tune#wandb_mixin
• https://medium.com/distributed-computing-with-ray/hyperparameter-optimization-for-transformers-a-guide-c4e32c6c989b
• https://colab.research.google.com/drive/1tQgAKgcKQzheoh503OzhS4N9NtfFgmjF?usp=sharing#scrollTo=ip6shHkNICTs