[TIL] WanDB 시작하기

예전에 다른 회사 면접을 봤을 때, 면접 질문 중 하나가 'WanDB 사용 해 봤어요?' 라는 것이었다. 말만 들어봤지 딱히 사용은 안 해봐서... 프로젝트에 들어가면 반드시 WanDB로 학습하고 있는 것 모니터링을 해보자 결심을 하였다.

WanDB는 모델을 학습하면서 데이터셋과 파라미터들을 관리하는 대표적인 MLOps 솔루션이다. 하나의 모델을 파인튜닝을 하다보면 어떤 모델이 최고인지 잊어버리기도 하고 결국 엑셀파일이나 로그파일을 뒤져가면서 정리할 때가 있는데 WanDB는 간단한 명령어를 통해서, 모델의 최적 파라미터들을 관리할 수 있다.

 

WanDB 가격 정보

 

WanDB를 로컬에서 사용할 것이라면, 간단하게 회원가입을 진행한 후 해당 터미널에서 아래와 같이 명령어를 입력하고 나서, 로그인후 보이는 API Key를 붙여놓으면 된다

pip install wandb
wandb login

API 키 붙여놓기

 

이러면 이제 WanDB 사용의 30%는 완료한 것이다. 이제 학습을 위해 열심히 파라미터를 조정하면 된다. 먼저 모델들의 train.py 파일들을 열어 wandb config들을 설정한 뒤 학습을 시키면 된다.

아마 많이 보는 train.py 파일들의 형태는 ArgumentParser 입력 받아서 파라미터들을 조정해줄텐데, 이때 아래와 같이 wandb.init(project='프로젝트명', configs=args)라는 한 줄을 넣고 시작을 하면 된다.

# config_experiment.py
import wandb
import argparse
import numpy as np
import random


# Training and evaluation demo code
def train_one_epoch(epoch, lr, bs):
    acc = 0.25 + ((epoch / 30) + (random.random() / 10))
    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))
    return acc, loss


def evaluate_one_epoch(epoch):
    acc = 0.1 + ((epoch / 20) + (random.random() / 10))
    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))
    return acc, loss


def main(args):
    # Start a W&B Run
    run = wandb.init(project="config_example", config=args)

    # Access values from config dictionary and store them
    # into variables for readability
    lr = wandb.config["learning_rate"]
    bs = wandb.config["batch_size"]
    epochs = wandb.config["epochs"]

    # Simulate training and logging values to W&B
    for epoch in np.arange(1, epochs):
        train_acc, train_loss = train_one_epoch(epoch, lr, bs)
        val_acc, val_loss = evaluate_one_epoch(epoch)

        wandb.log(
            {
                "epoch": epoch,
                "train_acc": train_acc,
                "train_loss": train_loss,
                "val_acc": val_acc,
                "val_loss": val_loss,
            }
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument("-b", "--batch_size", type=int, default=32, help="Batch size")
    parser.add_argument(
        "-e", "--epochs", type=int, default=50, help="Number of training epochs"
    )
    parser.add_argument(
        "-lr", "--learning_rate", type=int, default=0.001, help="Learning rate"
    )

    args = parser.parse_args()
    main(args)

 

또한 WanDB에서는 CPU와 GPU, 네트워크 상용량에 되해서도 확인해줄 수 있다.

 

 

참고 : WanDB 공식 문서 https://docs.wandb.ai/quickstart

Quickstart | Weights & Biases Documentation