[Wandb] Wandb 사용법 + Sweep with Pytorch

Contents

 

Weight & Biases

Wandb (Weight & Biases)는 MLOps 플랫폼이며, 딥러닝을 학습하기 위해 필요한 기능들을 제공한다. 

 

 

 

wandb.ai 에서 회원가입을 하게 되면 API 키를 발급받을 수 있다. 해당 API를 통해 wandb에 login하면 사용할 수 있다.

pip install wandb
wandb.login()

 

 

Wandb init

Wandb를 사용하기 위해서는 처음에 init을 설정해주어야 한다. 기본적으로 하이퍼 파라미터(Hyper-Paramters)들은 Yaml로 관리하는 것이 편하기 때문에 아래와 같이 사용했다. 

learning_rate: 1e-5
batch_size: 64
epochs: 10

import yaml 

def load_yaml(path):
    with open(path) as f:
        config = yaml.safe_load(f)
    return config 
    
args = load_yaml('config.yaml')

wandb.init(
project="project_name", 
name=f'{args.batch_size}-{args.lr}.pt',
config={
    'learning_rate': args.lr, 
    'batch_size': args.batch_size,
    'model_name': args.model_name, 

})

# Update
# wandb.config.update({'dr_rate':0.1})

# when using argparse
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', default=32)
parser.add_argument('--wandb', action='store_true')
parser.add_argument('--lr', default=1e-5)
parser.add_argument('--num_epochs', default=20)
parser.add_argument('--max_length', default=160)
parser.add_argument('--patient', default=20)
parser.add_argument('--model_name', default='bert-base-uncased')
parser.add_argument('--device', default='cuda:0')
args = parser.parse_args()
wandb.config.update(args)

 

 

Wandb를 통한 감성분석

데이터셋은 아무 감성분석 데이터를 가지고 오면 된다. 아래와 같이 코드를 작성할 수 있다.

 

Build DataLoader

먼저 입력으로 들어오는 Data를 Pytorch 입력으로 사용하기 위해 Dataset과 DataLoader를 구성하여야 한다. 

import os 
import pandas as pd 

import torch 
from torch.utils.data import Dataset, DataLoader

from transformers import AutoTokenizer

from settings import * 

class SentenceDataset(Dataset):
    def __init__(self, X, y=[], max_length=160, model_name='bert-base-uncased'):
        self.max_length = max_length
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, max_length=max_length)
        self.sent_cols = ['sentence_1', 'sentence_2']
        self.target_col = ['label']
        
        # require modification
        self.X = self.df2list(X.loc[:, 'sentence_1'])
        self.y = [i/5 for i in self.df2list(y) if self.df2list(y)]
    
    def __getitem__(self, idx):
        X = self.X[idx]
        X = self.tokenizer(
            X,
            add_special_tokens=True, 
            max_length=self.max_length, 
            padding='max_length',
            truncation=True
        )
        if self.y:
            return (
                {'input_ids': torch.tensor(X['input_ids'], dtype=torch.long), 
                 'type_ids': torch.tensor(X['token_type_ids'], dtype=torch.long), 
                 'mask': torch.tensor(X['attention_mask'], dtype=torch.long)}, 
                torch.tensor(self.y[idx], dtype=torch.float)
            )
        else:
            return {'input_ids': torch.tensor(X['input_ids'], dtype=torch.long), 
                    'type_ids': torch.tensor(X['token_type_ids'], dtype=torch.long), 
                    'mask': torch.tensor(X['attention_mask'], dtype=torch.long)}
    
    def __len__(self):
        return len(self.X)
    
    def df2list(self, y):
        try:
            return y.values.tolist()
        except:
            return []
        

def get_trainloader(args):
    path = os.path.join(DATA_DIR, 'train.csv')
    dataframe = pd.read_csv(path)
    X, y = dataframe, dataframe.loc[:, 'label']
    d_set = SentenceDataset(X, y, max_length=args.max_length, model_name=args.model_name)
    return DataLoader(d_set, batch_size=args.batch_size, shuffle=True)

def get_validloader(args):
    path = os.path.join(DATA_DIR, 'dev.csv')
    dataframe = pd.read_csv(path)
    X, y = dataframe, dataframe.loc[:, 'label']
    d_set = SentenceDataset(X, y,max_length=args.max_length, model_name=args.model_name)
    return DataLoader(d_set, batch_size=args.batch_size, shuffle=False)

def get_testloader(args):
    path = os.path.join(DATA_DIR, 'test.csv')
    X = pd.read_csv(path)
    d_set = SentenceDataset(X, max_length=args.max_length, model_name=args.model_name)
    return DataLoader(d_set, batch_size=args.batch_size, shuffle=False)

 

`DATA_DIR`로 경로를 잡아주고, self.X 와 self.y는 입력 DataFrame에 맞는 Column으로 변경해주면 된다. 기본적인 경로는 `settings.py` 파일에 적재해두고 import 해서 사용한다.

import os 

BASE_DIR = os.path.dirname(__file__)
DATA_DIR = os.path.join(BASE_DIR, 'data')
PARAM_DIR = os.path.join(BASE_DIR, 'parameters')
OUT_DIR = os.path.join(BASE_DIR, 'results')

 

 

Build Model

모델은 목적에 맞게 불러오면 된다. 만약에 Classification 단계를 바꾸거나 변형을 주고 싶은 경우에는 `AutoModel`을 불러와서 사용하면 되고, Classification이나 Regression을 사용하기 위해서는 아래와 같이 변경하면 된다.

from torch import nn 
from transformers import AutoModelForSequenceClassification, AutoModel

class Model(nn.Module):
    def __init__(self, model_name='jhgan/ko-sroberta-sts', input_dim=768, num_labels=1):
        super(Model, self).__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
        self.num_labels = num_labels
        self.input_dim = input_dim
        
        self._init_weight()
    
    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                nn.init.uniform_(m.bias)
    
    def forward(self, input_ids, type_ids, mask):
        pooler = self.model(input_ids=input_ids, token_type_ids=type_ids, attention_mask=mask)
        logits = pooler['logits']
        return logits

 

Classification Task를 다루기 위해서는 마지막 num_labels=2로 지정하여 마지막에 Sigmoid 함수를 사용하는 등으로 변경하면 되고, Regression Task를 다루기 위해서는 위와 같이 num_labels=1로 지정해서 사용하면 된다.

 

BERT의 기본 구조는 input_ids, token_type_ids, attention_mask를 입력으로 받기 때문에 이전 DataLoader에서 받은 input_ids, token_type_ids, attention_mask 를 사용하면 된다.

 

Build Function

추가적으로 분석에 필요한 다른 함수들은 아래와 같이 정의했다.

def elapsed_time(start, end):
    elapsed = end - start 
    elapsed_min = elapsed // 60 
    elapsed_sec = elapsed - elapsed_min * 60 
    return elapsed_min, elapsed_sec 

def torch2npy(tensor):
    if len(tensor.shape) == 4:
        tensor = tensor.unsqueeze(0)
    npy = tensor.detach().cpu().numpy()
    return npy

def log(args, epoch, elapsed_min, elapsed_sec, train_loss, valid_loss):
    print(f'Epoch: [{epoch+1}/{args.num_epochs}]\tElapsed Time: {elapsed_min}m {elapsed_sec:.2f}s')
    print(f'Train Loss: {train_loss:.4f}\tValid Loss: {valid_loss:.4f}')
    
    
def evaluate(args, model, valid_loader, criterion):
    valid_loss = 0
    with torch.no_grad():
        model.eval()
        for batch in valid_loader:
            X = {'input_ids': batch[0]['input_ids'].to(args.device), 
                 'mask': batch[0]['mask'].to(args.device)}
            y = batch[1].to(args.device)
            
            pred_y = model(**X).squeeze()
            loss = criterion(pred_y, target=y)
            valid_loss += loss.item()
        valid_loss /= len(valid_loader)
    return valid_loss
    
    
    
def predict(args, model, test_loader):
    pred_ys = []
    with torch.no_grad():
        model.eval()
        for batch in test_loader:
            X = {'input_ids': batch['input_ids'].to(args.device), 
                 'mask': batch['mask'].to(args.device)}
            
            pred_y = model(**X).squeeze()
            pred_ys.append(pred_y)
    return list(round(float(i)*5, 1) for i in torch.cat(pred_ys))

 

 

 

Sweep 사용하기

Sweep을 사용하기 위해서는 Paramter를 먼저 설정해주어야 한다. 위 사용한 예시와 마찬가지로 yaml 파일을 사용할 수도 있고, 아래와 같이 Dictionary를 사용할 수도 있다.

sweep_config = {
    'method': 'bayes'
}
metric = {
    'name': 'val_loss', 
    'goal': 'minimize'
}
parameters_dict = {
    'learning_rate': {
        'min': 1e-8,
        'max': 5e-4 
    }, 
    'epochs':{
        'value': 10
    }, 
    'batch_size': {
        'value': 64
    }
}


sweep_config['parameters'] = parameters_dict
sweep_config['metric'] = metric

 

 

 

Trainer

마지막 단계인 Training 단계를 구축할 차례다. 만약 추가적으로 Optimizer나 Loss Function을 수정해야 되는 경우에는 새로운 함수를 만들어서 'Adam' 이 입력으로 들어오면 Adam Optimizer, 'sgd'가 입력으로 들어오면 SGD Optimizer를 사용하는 형태로 코드를 작성할 수도 있다. 

def trainer(config=None):
    train_losses, valid_losses = [], []
    best_loss = float('inf')
    count = 0
    model = Model().to(args.device)
    
    train_loader = get_trainloader(args)
    valid_loader = get_validloader(args)
    criterion = nn.L1Loss()
    
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay_rate': 0.1},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay_rate': 0.0}
    ]
    
    with wandb.init(config=config) as run:
        config = wandb.config
        run.name = f'{args.model_name.split("/")[-1]}-{args.batch_size}-{config.learning_rate}.pt'
        optimizer = optim.AdamW(optimizer_grouped_parameters, lr=config.learning_rate)
        
        for epoch in tqdm(range(config.epochs), total=config.epochs):
            train_loss = 0 
            model.train()
            start = time.time()
            for batch in train_loader:
                optimizer.zero_grad()
                X = {'input_ids': batch[0]['input_ids'].to(args.device), 
                    'mask': batch[0]['mask'].to(args.device)}
                y = batch[1].to(args.device)

                pred_y = model(**X).squeeze()
                loss = criterion(pred_y, target=y)
                loss.backward()
                train_loss += loss.item()
                optimizer.step()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            train_loss /= len(train_loader)
            valid_loss = evaluate(args, model, valid_loader, criterion)
            end = time.time()
            elapsed_min, elapsed_sec = elapsed_time(start, end)
            
            log(args, epoch, elapsed_min, elapsed_sec, train_loss, valid_loss)
            if args.wandb:
                wandb.log({
                    'train_loss': train_loss, 
                    'val_loss': valid_loss
                })
            
            if best_loss > valid_loss:
                best_loss = valid_loss
                torch.save(model.state_dict(), os.path.join(PARAM_DIR, f'{args.model_name.split("/")[-1]}-{args.batch_size}-{config.learning_rate}.pt'))
                count = 0
            else:
                count += 1
                if count == args.patient:
                    print(f'Early Stopping! [Patient={args.patient}]')
                    break 
            
            train_losses.append(train_loss)
            valid_losses.append(valid_loss)
    return train_losses, valid_losses

 

 

 

Run

아래 코드 한줄로 몇 번 학습할 것인지에 대한 결과를 출력할 수도 있다.

wandb.agent(sweep_id, trainer, count=20)