[MLOps] 1. WandB 입문
1. WandB 란?
Weight and Bias 의 줄임말로 머신러닝 / 딥러닝 실험과정을 쉽게 관리할 수 있도록 도와주는 tool 이다. 딥러닝 연구를 하다보면 정말 많은 모델을 다루게 되고 모델에 크고 작은 변화를 주게 된다. 모델 구조의 변화 뿐만 아니라 하이퍼 파라미터, 데이터 전처리 등 여러가지 요인으로 인해서 수십 ~ 수백개의 경우의 수를 모두 비교해야하는 상황이 생긴다. 이때 보통 jupyter notebook 으로 실험결과를 시각화, 비교하는 경우가 많았는데 실험이 많아질수록 이 방법에는 한계가 존재하게 된다. 나는 보통 엑셀을 통해 실험결과를 정리하고 하이퍼 파라미터 튜닝, 시각화 코드를 직접 작성해서 일일이 비교를 진행했는데 보통 한달에 한번씩은 꼭 결과들이 꼬여서 다시 정리하는 일이 생겼고 그럴때마다 지옥을 경험했다... 지금까지는 이렇게 하는게 실험과 모델을 이해하는데 더 도움이 될 것이라고 판단했지만 그냥 미련했던 것 같다. WandB는 이러한 과정을 정말 보기 쉽고 편하게 정리해준다. 이러한 tool 들은 보통 MLOps 와 연관지어서 언급이 많이 되는데 WandB, MLflow, Kubernetes 등 다양한 툴이 존재하며 각자의 장/단점들이 명확하게 존재하기 때문에 본인의 용도에 맞춰서 사용하도록 하자. WandB는 개인적인 용도로 실험만을 관리하기에 가장 적합해보였다.
2. WandB 시작하기
우선 wandb 홈페이지에 가서 가입을 한다. 구글, 깃허브 아이디를 사용하면 쉽게 가입할 수 있고 기관을 적으라는데... 나는 그냥 아무거나 적어서 제출했다. 가입 후 로그인을 하면 매우 친절하게 가이드를 알려주기 때문에 어렵지 않게 시작할 수 있다. 나는 jupyter notebook 을 통해 작업했는데 notebook 에서 로그인을 했더니 무언가 오류가 발생해서 그냥 pychram 터미털에서 작업을 해주었다. 어차피 코드 관리를 위해서는 pycharm 으로 연습하는게 좋아보인다.
3. 소스코드
모든 코드는 깃허브에 업로드하였고 아래에서는 주요 코드만 설명했다. 학습 관련 utils들은 링크를 통해 확인할 수 있다.
https://github.com/Hyunmok-Park/WandB_practice
GitHub - Hyunmok-Park/WandB_practice
Contribute to Hyunmok-Park/WandB_practice development by creating an account on GitHub.
github.com
3.1 모델(net.py)
모델은 아주 간단한 CNN을 사용했다. 코드는 pytorch 를 사용해서 작성하였고, 데이터는 torchvision 에서 제공하는 mnist 데이터를 사용했다.
from torch import nn
class ConvNet(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 32, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(2, 2)
)
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, 3, 1, 1),
nn.ReLU(),
nn.MaxPool2d(2, 2)
)
self.layer3 = nn.Sequential(
nn.Linear(64 * 7 * 7, 128, bias=True),
nn.ReLU()
)
self.layer4 = nn.Sequential(
nn.Linear(128, 84),
nn.ReLU()
)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = x.view(x.size(0), -1)
x = self.layer3(x)
x = self.layer4(x)
x = self.fc3(x)
return x3.2 학습함수(train_utils.py)
모델의 파라미터를 업데이트하고, wandb 라이브러리의 watch, log 함수를 사용해 학습 과정에 대한 로그를 기록할 수 있다.
import numpy as np
def train(model, train_loader, criterion, optimizer, device, config, wandb, epoch):
model.train()
wandb.watch(model, criterion, log="all", log_freq=10)
loss_list = []
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
output = model(images)
loss = criterion(output, labels)
loss_list += [loss.detach().cpu().numpy().tolist()]
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss = np.mean(loss_list)
wandb.log({"train_loss": avg_loss}, step=epoch)
print(f"TRAIN: EPOCH {epoch + 1:04d} / {config.epochs:04d} | Epoch LOSS {avg_loss:.4f}")
def vaild(model, vali_loader, criterion, device, wandb, epoch):
model.eval()
val_loss = []
for data, target in vali_loader:
data, target = data.to(device), target.to(device)
output = model(data)
val_loss += [criterion(output, target).item()]
val_loss = np.mean(val_loss)
wandb.log({"valid_loss": val_loss}, step=epoch)
print(f"VALID: LOSS {val_loss:.4f} | Accuracy {val_loss:.4f} ")3.3 메인함수(main.py)
WandB의 과정을 시작할 수 있는 코드로 wandb 라이브러리의 sweep, agent, init 함수를 사용해서 모든 과정을 시작할 수 있다. 이때 sweep 이라는 기능이 가장 중요한 기능인데 hyperparameter tuning 을 자동화할 수 있는 기능이다. WandB는 기본적으로 yaml 형식의 입력을 통해 관리된다. wandb.init() 함수에 우리가 원하는 모델의 스펙을 dictionary 형태(yaml)로 넣어주면 해당 스펙을 일종의 primary key 로서 활용하여 기록을 시작한다. 그러면 wandb 페이지에서는 run의 형태(아래의 run 함수와는 다름)로 해당 실험을 기록한다. run과 init이 1대1 형태로 기록이 되는 것이다.(아직 입문자라서 이 부분은 확실하지 않음...)
그리고 sweep 이라는 기능은 run-init 쌍을 우리가 원하는 만큼 자동으로 돌려주는 것이다. sweep도 init 과 마찬가지로 dictionary 형태로 관리된다. 다만 init 에 들어가는 형태와는 다르게 각 하이퍼 파라미터의 범위를 지정하게 된다. 그러면 sweep 함수는 알아서 init 에 들어가게 될 하나의 config를 만들어주고 run 을 형성한다. 우리는 sweep, agent 함수를 사용해서 sweep 전용 config, init 전용 config를 넣어줄 함수(아래는 run) 을 지정해주면 알아서 모든 것을 해준다.
import torch
from torch import nn
import wandb
from tqdm import tqdm
from dataloader import Load_Dataset
from train_utils import train, vaild
from net import ConvNet
from utils import build_optimizer
from make_config import *
DEVICE = torch.device('cuda')
def run(config=None):
wandb.init(config=config)
w_config = wandb.config
criterion = nn.CrossEntropyLoss()
train_loader, vaild_loader = Load_Dataset(w_config.batch_size)
model = ConvNet().to(DEVICE)
optimizer = build_optimizer(model, w_config.optimizer, w_config.learning_rate)
for epoch in tqdm(range(w_config.epochs), desc='EPOCH'):
train(model, train_loader, criterion, optimizer, DEVICE, w_config, wandb, epoch)
vaild(model, vaild_loader, criterion, DEVICE, wandb, epoch)
if __name__ == '__main__':
sweep_config = sweep_config
sweep_id = wandb.sweep(sweep_config, project="sweep_tutorial", entity='hmpark1995')
wandb.agent(sweep_id, run)
# default_config = hyperparameter_defaults
# run(default_config)3.3 config(make_config.py)
sweep, init 에 넣어줄 config를 미리 지정해둔 함수이다. 이번에는 4개의 하이퍼 파라미터를 조정했다. values 라는 이름으로 원하는 파라미터 집합을 넣어주면 모든 조합을 알아서 측정해준다. 각 key 가 어떤 동작을 해주는 것인지는 점차 알아가도록 하자.
sweep_config = {
'name' : 'bayes-test',
'method': 'random',
'metric' : {
'name': 'valid_loss',
'goal': 'minimize'
},
'parameters' : {
'optimizer': {
'values': ['adam', 'sgd']
},
'epochs': {
'values': [3, 4]
},
'learning_rate': {
'values': [0.1, 0.01]
},
'batch_size': {
'values': [32, 64]
}
}
}4. 결과
단일 실험(init)을 진행했을때 결과. 내가 log에 기록했던 valid, train loss 가 step에 맞추어 기록되었으며 watch 함수를 사용하면 아래처럼 gradients, parmeters도 확인할 수 있다.
Sweep 을 통한 하이퍼 파라미터 튜닝 결과. log에 기록한 loss는 여기서도 확인이 가능하며 sweep 을 통해 최소화하고자 했던 valid_loss에 대한 기록은 더 자세하게 확인이 가능하다. 그리고 놀랍게도 각 하이퍼 파라미터와 sweep 목표의 연관성을 자동으로 측정해준다...
5. 요약
이번 포스트에서는 MLOps 에서 유명한 tool 중 하나인 WandB를 입문해보았다. 내가 생각했던 기능보다도 더 많은 기능을 제공했고 사용하기도 매우 쉬웠다. 다른 MLOps tool들과 비교해 serving 에는 조금 취약하다고는 하지만 개인적인 용도로서 실험을 관리하고 추적하는 목적에서는 매우 좋은 도구라고 생각된다.