[wandb] wandb에 대해 알아보고 사용해보자

 

머신러닝 및 딥러닝 프로젝트에서 모델의 추적을 위해 강력한 도구인 wandb(Weights and Biases)를 사용하여 모델의 성능을 극대화하는 방법을 소개하려 한다.

 

wandb는 무엇이고 왜 사용하는가

wandb(Weights and Biases)는 머신러닝 및 딥러닝 프로젝트에서 모델의 파라미터에 대한 도구로, 실험 결과를 시각화하고, 하이퍼파라미터를 조정하며, 모델의 성능을 모니터링하는 도구이다.   특히, 하이퍼파라미터 최적화를 통해 최적의 모델 설정을 찾는 데 큰 도움이 된다.

 

wandb 가입 및 로그인

wandb를 사용하기 위해서는 wandb 공식 웹사이트에서 회원가입과 로그인을 해야 합니다. 이메일 주소와 비밀번호를 입력하거나, 구글 계정이나 GitHub 계정으로 손쉽게 가입하고 로그인할 수 있다.

가입버튼 클릭

공식사이트에 들어가서 가입버튼을 클릭한다.

나의 경우는 Github로 가입하고 개인정보를 입력한다.

다음과 같은 창이 나오게 되는데 여기서 API KEY를 복사해 두고 저장해 둔다.

 

wandb 로그인

pip install wandb

wandb 설치

import wandb
wandb.login()

다음과 같이 셀을 실행하면 밑에와 같은 창이 생긴다. 이때, 여기에 복사한 API KEY를 입력한다.

 

API KEY를 입력하면 다음과 같이 True가 출력된다.

혹시 키를 까먹었다면 여기를 클릭하면 된다.

 

wandb를 사용하는 방법

훈련코드

import pytorch_lightning as pl
import wandb
from pytorch_lightning.loggers import WandbLogger
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import transforms
from torchvision.datasets import MNIST

# 1. wandb Logger 설정
wandb_logger = WandbLogger(project="pytorch-lightning-demo")

# 2. 데이터 준비
class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, data_dir='./', batch_size=32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def prepare_data(self):
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        transform = transforms.Compose([transforms.ToTensor()])
        self.mnist_train = MNIST(self.data_dir, train=True, transform=transform)
        self.mnist_val = MNIST(self.data_dir, train=False, transform=transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)

# 3. 모델 정의
class LitMNIST(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 256)
        self.layer_3 = torch.nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.layer_1(x))
        x = F.relu(self.layer_2(x))
        x = self.layer_3(x)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        self.log('val_loss', loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

# 4. 학습
if __name__ == '__main__':
    dm = MNISTDataModule()
    model = LitMNIST()
    trainer = pl.Trainer(logger=wandb_logger, max_epochs=5)
    trainer.fit(model, dm)

위의 코드를 요약하면 다음과 같습니다. 

import wandb
from pytorch_lightning.loggers import WandbLogger

wandb_logger = WandbLogger(project="pytorch-lightning-demo")

trainer = pl.Trainer(logger=wandb_logger, max_epochs=5)

WandbLogger를 선언한 뒤, project 이름을 작성해 준 뒤 trainer 인수로 넣어주고 trinaing을 진행하면 wandb 사이트에 훈련하면서 loss값이 변하는 것을 확인할 수 있다. 

 

wandb를 통해 빠르게 모델 최적화를 하시기를 바랍니다!