NLP - Pytorch Finetune LightningModule

이번 포스팅에서는 PyTorch에 대한 High-level 인터페이스를 제공하는 오픈소스 Python 라이브러인 Lightning에 대해서 다뤄보려고 합니다. Lightning이 등장한 배경에 대해서 알아보고 직접 구현하는 작업까지 진행해보겠습니다. 구현은 KcELECTRA의  NCMC downstream task를 통해서 구현하고 이해해보는 과정을 진행할 것입니다. 실제로 KcELECTRA에서 downstream 예시코드에서도 사용되었고 다양한 한국어 자연어 모델의 Finetune코드에서 사용되고 있는 라이브러리 입니다. 

 

 KcELECTRA 모델의 제작자분께서 공유해주신 Fine-tune 코드 링크입니다.

https://colab.research.google.com/drive/11WQdZSf_1xIcRrcHRmqXI4sbowF27G-o#scrollTo=Yb0113DUFE1k

Google Colaboratory Notebook

 

LightningModule 등장배경

---

LightningModule은 torch를 통해 모델을 만드는데 이전에는 각각의 train, validation, test, loss, optimzer등 모델 내부에서 작동하는 기능들을 각각 따로 구현하여 작동시켜야 했습니다. 그런데 GANs과 같이 두 개의 모델을 학습시켜야 하는 모델의 경우 그 구조가 굉장히 복잡해집니다. 또한, 이렇게 구현하게 되면 각 모델을 구성한 제작자의 스타일에따라 다른 형태로 구현되어 서로 정보에 대한 공유가 어려운점 역시 존재합니다. 그래서 등장한 것이 LightningModule입니다. LightningModule 사용하면 하나의 클래스에서 따로 구성해야했던 부분을 수정, 추가를 진행할 수 있습니다.

 

 

LightningModule 기본구조

---

LightningMod을 통해 모델을 구성하는 기본틀은 다음과 같습니다. 아래 4개의 메서드 중 forward를 제외하고 꼭 구현해줘야하는 메서드들입니다. 그런데 forward에 대한 구현을 하는게 좋습니다. 해당 메서드를 정의해두면 클래스 내부에서 순전파에 대한 기능을 self(<값>)으로 편하게 사용이 가능합니다. 

import pytorch_lightning as pl

class LitModel(pl.LightningModule):
    def __init__(self, vocab_size):
        super().__init__()
        self.model = Transformer(vocab_size=vocab_size)
        
    def forward(self, inputs, target):
        return self.model(inputs, target)

    def training_step(self, batch, batch_idx):
        inputs, target = batch
        output = self(inputs, target)
        loss = torch.nn.functional.nll_loss(output, target.view(-1))
        return loss

    def configure_optimizers(self):
        return torch.optim.SGD(self.model.parameters(), lr=0.1)

위의 구조는 다음과 같은 형태로 내부적으로 작동합니다.

model.train()
torch.set_grad_enabled(True)

for batch_idx, batch in enumerate(train_dataloader):
    loss = training_step(batch, batch_idx)

    # clear gradients
    optimizer.zero_grad()

    # backward
    loss.backward()

    # update parameters
    optimizer.step()

 

LightningModule 구현해보기(KcELECTRA - NSMC Finetune)

---

코드는 위에 링크에 있는 코드를 최대한 건드리지 않고 진행하였습니다. 구현해보기의 순서는 파라미터 초기화 -> 분류를 위한 모델 구조 이해하기 순으로 진행해보겠습니다.

 

1) 파라미터(Parameter)

우선 모델을 불러와 설정하기 위한 파라미터에 대한 내용입니다. 여기서 정의한 파라미터들은 모델 클래스에 전달되어 사용될 파라미터들의 값을 정의한 것입니다.

args = {
    'random_seed': 42, # Random Seed
    'pretrained_model': 'beomi/KcELECTRA-base',  # Transformers PLM name
    'pretrained_tokenizer': '',  # Optional, Transformers Tokenizer Name. Overrides `pretrained_model`
    'batch_size': 32,
    'lr': 5e-6,  # Starting Learning Rate
    'epochs': 1,  # Max Epochs
    'max_length': 150,  # Max Length input size
    'train_data_path': "nsmc/ratings_train.txt",  # Train Dataset file
    'val_data_path': "nsmc/ratings_test.txt",  # Validation Dataset file
    'test_mode': False,  # Test Mode enables `fast_dev_run`
    'optimizer': 'AdamW',  # AdamW vs AdamP
    'lr_scheduler': 'exp',  # ExponentialLR vs CosineAnnealingWarmRestarts
    'fp16': True,  # Enable train on FP16(if GPU)
    'tpu_cores': 0,  # Enable TPU with 1 core or 8 cores
    'cpu_workers': os.cpu_count(),
}

 

 

2) LightningModule

2-1) __init__

init의 경우 클래스의 경우 주의해서 확인해 봐야할 부분은 처음 호출되는 save_hyperparameters()입니다. 해당 메서드는 하이퍼파라미터에 대한 정보를 갖는 체크포인트를 만들고 'self.hparams'에 정보가 저장됩니다. 저장되는 방식에 대한 설명은 코드에 추가했듯 해당 메서드를 호출하면 모든 매개변수를 self.hparams에 저장합니다. 특정 파라미터들만 저장하고 싶다면 지정해주면 됩니다. 이후에는 감성분류를 위한 모델과 토크나이저를 불러오는 부분입니다. save_hyperparameters를 호출했기 때문에 위에서 정의한 파라미터들을 사용할 수 있습니다.

from pytorch_lightning import LightningModule, Trainer, seed_everything
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AdamW

class Model(LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        self.save_hyperparameters() # 해당 메서드는 init의 self 이후에 오는 모든 파라미터를 self.hparams 저장합니다. 
        """
        def __init__(self, param1 = 128, param2 = 64)로 설정했다면 
        self.save_hyperparameters("param1", "param2")를 통해서도 self.hparams에 저장할 수 있습니다.
        """

        self.clsfier = AutoModelForSequenceClassification.from_pretrained(self.hparams.pretrained_model)
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.hparams.pretrained_tokenizer
            if self.hparams.pretrained_tokenizer
            else self.hparams.pretrained_model
        )

 

2-2) forword

순전파 역할을 합니다. 해당 메서드를 정의했다면 간편하게 self(<값>)으로 클래스 내부에서 해당 메서드를 호출합니다.

def forward(self, **kwargs):
        return self.clsfier(**kwargs)

 

2-3) step

앞에서 Lightning구조에서 없던 부분입니다. traing_step과 validation_step은 라이브러리 안에 정의 되어 있는 부분인데 동일한 과정을 진행하므로 해당 과정에 대해서 사용자 메서드를 만들고 각각의 메서드에서 step메서드를 호출하여 코드를 깔끔하게 작성하였습니다. 결과(output)을 받는 부분을 보시면 self(input_ids=data, labels=labels)로  forward를 간편하게 호출한 모습을 확인하실 수 있습니다. 결과값과 loss를 딕셔너리 형태로 반환합니다. 딕셔너리 형태로 반환할때는 'loss' key는 반드시 들어가야합니다. 혹은 loss tensor만 반환해도 됩니다.

def step(self, batch, batch_idx):
        data, labels = batch
        output = self(input_ids=data, labels=labels)

        # Transformers 4.0.0+
        loss = output.loss
        logits = output.logits

        preds = logits.argmax(dim=-1)

        y_true = list(labels.cpu().numpy())
        y_pred = list(preds.cpu().numpy())

        return {
            'loss': loss,
            'y_true': y_true,
            'y_pred': y_pred,
        }

def training_step(self, batch, batch_idx):
        return self.step(batch, batch_idx)

def validation_step(self, batch, batch_idx):
        return self.step(batch, batch_idx)

 

2-4) epoch_end

on_train_epoch_end, on_validation_epoch_end, on_test_epoch_end가 존재하는데 이 메서드들은 훈련, 검증, 테스트 과정중1epoch가 끝났을때 각각의 step에서 호출합니다. 링크에 있는 코드에서는 train_epoch_end로 되어 있지만 라이브러리 업데이트로 인해 메서드의 이름과 구조가 변경 되었습니다.

def epoch_end(self, outputs, state='train'):
        loss = torch.tensor(0, dtype=torch.float)
        for i in outputs:
            loss += i['loss'].cpu().detach()
        loss = loss / len(outputs)

        y_true = []
        y_pred = []
        for i in outputs:
            y_true += i['y_true']
            y_pred += i['y_pred']
        
        acc = accuracy_score(y_true, y_pred)
        prec = precision_score(y_true, y_pred)
        rec = recall_score(y_true, y_pred)
        f1 = f1_score(y_true, y_pred)

        self.log(state+'_loss', float(loss), on_epoch=True, prog_bar=True)
        self.log(state+'_acc', acc, on_epoch=True, prog_bar=True)
        self.log(state+'_precision', prec, on_epoch=True, prog_bar=True)
        self.log(state+'_recall', rec, on_epoch=True, prog_bar=True)
        self.log(state+'_f1', f1, on_epoch=True, prog_bar=True)
        print(f'[Epoch {self.trainer.current_epoch} {state.upper()}] Loss: {loss}, Acc: {acc}, Prec: {prec}, Rec: {rec}, F1: {f1}')
        return {'loss': loss}
    
def training_epoch_end(self, outputs):
    self.epoch_end(outputs, state='train')

def validation_epoch_end(self, outputs):
    self.epoch_end(outputs, state='val')

 

++ 위 코드는 구 버전으로 실행되지 않습니다. 아래 형식으로 수정하여 진행해야합니다. 직접 수정하여 진행해보시는걸 추천드립니다.

    def __init__(self):
        super().__init__()
        self.training_step_outputs = {'loss' : [], 'y_true' : [], 'y_pred' : []}

    def training_step(self):
        loss = ...  
        y_true = ...
        y_pred = ...
        
        self.training_step_outputs['loss'].append(loss)
        self.training_step_outputs['y_true'].append(y_true)
        self.training_step_outputs['y_pred'].append(y_pred)
        
        return loss

    def on_train_epoch_end(self, outputs):
        epoch_mean = torch.stack(outputs['loss']).mean()
        
        y_true = []
        y_pred = []
        for i in outputs:
            y_true += i['y_true']
            y_pred += i['y_pred']

 

2-5) configure_optimizers

옵티마이저를 설정하는 함수입니다. 모델 초기화시 하이퍼파라미터 값을 넘겨줬기 때문에 hparams값을 사용해 옵티마이저와 러닝레이트를 설정하고 스케줄러를 설정합니다.

def configure_optimizers(self):
        if self.hparams.optimizer == 'AdamW':
            optimizer = AdamW(self.parameters(), lr=self.hparams.lr)
        elif self.hparams.optimizer == 'AdamP':
            from adamp import AdamP
            optimizer = AdamP(self.parameters(), lr=self.hparams.lr)
        else:
            raise NotImplementedError('Only AdamW and AdamP is Supported!')
        if self.hparams.lr_scheduler == 'cos':
            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=1, T_mult=2)
        elif self.hparams.lr_scheduler == 'exp':
            scheduler = ExponentialLR(optimizer, gamma=0.5)
        else:
            raise NotImplementedError('Only cos and exp lr scheduler is Supported!')
        return {
            'optimizer': optimizer,
            'scheduler': scheduler,
        }

 

2-6) dataloader

실제 LightningModule 클래스에서 호출하는 함수는 train_dataloader, val_dataloader 두 함수입니다. dataloader는 역시 제작자분께서 데이터를 가져와 모델에 맞는 데이터셋을 만드는 과정을 함수화 하셨습니다. 코드 내용은 전부 적기에는 길어져서 생략하였습니다. 자세한 코드는 직접 colab링크를 확인해주세요.

def dataloader(self, path, shuffle=False):
    df = self.read_data(path)
    df = self.preprocess_dataframe(df)

    dataset = TensorDataset(
        torch.tensor(df['document'].to_list(), dtype=torch.long),
        torch.tensor(df['label'].to_list(), dtype=torch.long),
    )
    return DataLoader(
        dataset,
        batch_size=self.hparams.batch_size * 1 if not self.hparams.tpu_cores else self.hparams.tpu_cores,
        shuffle=shuffle,
        num_workers=self.hparams.cpu_workers,
    )

def train_dataloader(self):
    return self.dataloader(self.hparams.train_data_path, shuffle=True)

def val_dataloader(self):
    return self.dataloader(self.hparams.val_data_path, shuffle=False)


"""
밑에 함수들은 데이터로드 및 전처리를 위해 만들어진 사용자 함수입니다. 링크 코드를 참고해주세요.
"""

def read_data(self, path):
    ...

def clean(self, x):
    ...

def encode(self, x, **kwargs):
    ...

def preprocess_dataframe(self, df):
    ...

 

 

2-7) 결과 확인을 위한 코드

해당 부분은 훈련된 모델에 실제 데이터를 대입하여 결과를 받아보기 위해 사용되는 코드입니다. model(<값>)을 호출하게 되면 우리가 얻는것은 해당 입력 데이터에 대한 결과 로짓(logit)값입니다. 이를 분류 문제를 풀기 위해서는 Softmax로 결과를 통과 시켜주면 각 클래스에 대한 확률 값을 얻을 수 있고 이후에 argmax를 사용한다면 가장 높은 확률의 클래스 값을 얻을 수 있게 됩니다.

# 클래스 확률 분포
def infer(x):
    return torch.softmax(model(**model.tokenizer(x, return_tensors='pt')).logits, dim=-1)

---

+ 실제로 LightningModule의 실행 순서는 다음과 같습니다. 해당 라이브러리에 대한 정보는 직접 위 링크를 따라 들어가서 확인해보고 직접 메서드들을 오버라이딩해서 구현해보는걸 추천드립니다.

 

다음 포스팅에서는 LightningModule을 활용해서 두 개의 모델을 학습시키는 GANs를 구현해보겠습니다. GANs를 구현해보는 이유는 두 개의 모델을 훈련시켜야하는 상황에서 torch만으로의 작업은 매우 복잡하게 구성해야하는데 이 과정이 얼마나 간단해지는지 확인해보기 위함입니다.

def fit(self):
    if global_rank == 0:
        # prepare data is called on GLOBAL_ZERO only
        prepare_data()

    configure_callbacks()

    with parallel(devices):
        # devices can be GPUs, TPUs, ...
        train_on_device(model)


def train_on_device(model):
    # called PER DEVICE
    setup("fit")
    configure_optimizers()
    on_fit_start()

    # the sanity check runs here

    on_train_start()
    for epoch in epochs:
        fit_loop()
    on_train_end()

    on_fit_end()
    teardown("fit")


def fit_loop():
    model.train()
    torch.set_grad_enabled(True)

    on_train_epoch_start()

    for batch in train_dataloader():
        on_train_batch_start()

        on_before_batch_transfer()
        transfer_batch_to_device()
        on_after_batch_transfer()

        out = training_step()

        on_before_zero_grad()
        optimizer_zero_grad()

        on_before_backward()
        backward()
        on_after_backward()

        on_before_optimizer_step()
        configure_gradient_clipping()
        optimizer_step()

        on_train_batch_end(out, batch, batch_idx)

        if should_check_val:
            val_loop()

    on_train_epoch_end()


def val_loop():
    on_validation_model_eval()  # calls `model.eval()`
    torch.set_grad_enabled(False)

    on_validation_start()
    on_validation_epoch_start()

    for batch_idx, batch in enumerate(val_dataloader()):
        on_validation_batch_start(batch, batch_idx)

        batch = on_before_batch_transfer(batch)
        batch = transfer_batch_to_device(batch)
        batch = on_after_batch_transfer(batch)

        out = validation_step(batch, batch_idx)

        on_validation_batch_end(out, batch, batch_idx)

    on_validation_epoch_end()
    on_validation_end()

    # set up for train
    on_validation_model_train()  # calls `model.train()`
    torch.set_grad_enabled(True)