KoBART 분석해보기 (2)

저번에 이어서 KoBART 코드를 분석해보고자 하였다. 저번에는 데이터셋을 불러오는 파일을 분석해보았으니 이번에는 불러온 데이터셋을 학습시키는 코드를 지닌 train.py 파일과, 학습시킨 모델을 바탕으로 추론을 진행할 수 있게끔 해주는 infer.py 파일을 분석해보았다.

 

 

 

우선 train.py 파일이다. 아래와 같이 필요한 라이브러리들을 import 해준다.

 

 

 

import argparse
import logging
import os
import numpy as np
import pandas as pd
import pytorch_lightning as pl
import torch
from pytorch_lightning import loggers as pl_loggers
from torch.utils.data import DataLoader, Dataset
from dataset import KobartSummaryModule
from transformers import BartForConditionalGeneration, PreTrainedTokenizerFast
from transformers.optimization import AdamW, get_cosine_schedule_with_warmup

 

 

 

argparse는 저번에도 알아보았듯이 명령창에서 명령어로 Python 파일을 실행시킬 때 인자를 전달해줄 수 있게끔 하기 위해 import 한다.

 

 

근데 이번에는 저번에 쓰이지 않았던 logging 라이브러리가 import 되어있다. logging 라이브러리 또한 Python 내장 라이브러리로, 어떠한 형식으로든지 log를 출력하고 싶을 때 사용되는 라이브러리라고 한다. (그래서 주로 Python 서버에서 많이 사용된다고 한다.) 

 

 

이를 사용하기 위해서 우선 logger 객체를 생성해준다. 아래와 같이 작성해주면 된다.

 

 

 

logger = logging.getLogger()

 

 

 

그 후, 로그의 출력 기준을 설정해줘야 하는데 그 심각도의 수준은 CRITICAL > ERROR > WARNING > INFO > DEBUG 순으로 낮아진다. 현재 분석하고 있는 KoBART는 출력기준을 INFO로 잡아두고 있는데, 이는 실행하는 코드가 예상대로 잘 작동하는지 확인할 수 있게끔 해준다. 2 가지 정도만 더 알아보자면 DEBUG로 할 시에는 상세한 정보를 출력해주어 보통 문제를 진단할 때 사용되고, ERROR로 설정하고 사용하면 수행하지 못한 문제들을 출력해준다.

 

 

또, logger 객체에 로그 출력 형식을 아래와 같이 지정해주고 사용할 수 있다.

 

 

 

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

stream_handler = logging.StreamHandler()
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)

file_handler = logging.FileHandler('log.log')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

 

 

StreamHandler을 사용하여 출력구문에 표시하는 것도 가능하고, FileHandler을 사용하여 로그 파일에 로그를 저장하는 것 또한 가능하다.

 

 

 

다시 KoBART 코드로 돌아오자. 우선 아래와 같이 logging 객체를 지정해준다.

 

 

 

parser = argparse.ArgumentParser(description='KoBART Summarization')

parser.add_argument('--checkpoint_path',
                    type=str,
                    help='checkpoint path')

logger = logging.getLogger()
logger.setLevel(logging.INFO)

 

 

 

위 코드를 보게되면 logger 외에도 저번 시간에 알아본 ArgumentParser를 사용하여 인자를 전달해주는 역할을 하게끔 하고있다.

 

 

 

class ArgsBase():
    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)
        parser.add_argument('--train_file',
                            type=str,
                            default='data/train.tsv',
                            help='train file')

        parser.add_argument('--test_file',
                            type=str,
                            default='data/test.tsv',
                            help='test file')

        parser.add_argument('--batch_size',
                            type=int,
                            default=14,
                            help='')
        parser.add_argument('--max_len',
                            type=int,
                            default=512,
                            help='max seq len')
        return parser

 

 

 

train.py 파일에 첫 번째로 선언되어 있는 Class이다. Class 안에 @staticmethod를 달고 있는 function이 하나있다. (Class를 따로 불러오지 않고도 함수를 사용할 수 있음) 함수 내용을 보게되면 ArgumentParser을 받아서 추가적으로 add_argument를 진행해주는 것 같다.

 

 

 

class Base(pl.LightningModule):
    def __init__(self, hparams, trainer, **kwargs) -> None:
        super(Base, self).__init__()
        self.save_hyperparameters(hparams)
        self.trainer = trainer

    @staticmethod
    def add_model_specific_args(parent_parser):
        # add model specific args
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)

        parser.add_argument('--batch-size',
                            type=int,
                            default=14,
                            help='batch size for training (default: 96)')

        parser.add_argument('--lr',
                            type=float,
                            default=3e-5,
                            help='The initial learning rate')

        parser.add_argument('--warmup_ratio',
                            type=float,
                            default=0.1,
                            help='warmup ratio')

        parser.add_argument('--model_path',
                            type=str,
                            default=None,
                            help='kobart model path')
        return parser

 

 

 

그 다음에는 Base라는 이름의 Class가 선언되어 있다. 이 Class는 길이가 좀 길어서 반씩 나누어서 알아보도록 하겠다.

 

 

 

우선 해당 Class는 Pytorch Lighting 라이브러리의 LightningModule Class를 상속받고 있는 것을 확인할 수 있다. Pytorch Lightning에서는 trainer과 Model이 서로 상호작용을 할 수 있도록 pytorch nn.Module의 상위 Class인 LightningModule을 구현하게 된다고 한다. 기존 PyTorch에서 DataLoader, Mode, Optimizer, Training Loof 등을 전부 따로따로 코드로 구현해야 하는 반면에 LightningModule을 사용하면 모든 것을 한 번에 구현할 수 있다. PyTorch Lightning 공식 Document에 따르면 LightningModule은 6개의 Section으로 나뉜다고 한다.

 

 

 

• Computations (init).
• Train Loop (training_step)
• Validation Loop (validation_step)
• Test Loop (test_step)
• Prediction Loop (predict_step)
• Optimizers and LR Schedulers (configure_optimizers)

 

 

 

밑의 __init__ 함수 안을 보게되면 save_hyperparameters 함수를 사용하여 하이퍼파라미터들을 지정해주고 있는 것을 확인할 수 있으며(인자를 넣어주지 않으면 init의 인자들을 전부 저장), trainer을 init의 인자로 받아주어 self.trainer로 지정해주는 것도 확인할 수 있다. 또, 처음 __init__() 함수를 보면 뒤에 -> None 이라고 적혀있는 것을 볼 수 있는데 이는 해당 함수의 리턴 값을 명시해줄 때 사용되는 것이라고 한다. (-> str 이라고 적으면 return 값이 str인 것.) 

 

 

 

아래 코드는 WikiDocs "PyTorch 딥러닝 챗봇"에서 가져온 코드이다. 아래와 같이 간단하게 Class를 작성해줄 수 있는 것이다.

 

 

 

import torch
import pytorch_lightning as pl
from torch import Tensor, nn
from sklearn.datasets import load_boston
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import Dataset, DataLoader
from torch.nn import functional as F
import numpy as np

#Boston 집값 데이터를 읽어온다.
X, y = load_boston(return_X_y=True)

class SklearnDataset(Dataset):
    def __init__(self, X: np.ndarray, y: np.ndarray):
        super().__init__()
        scaler = MinMaxScaler() 

        scaler.fit(X) 
        self.X = scaler.transform(X)
        self.Y = y

    def __len__(self):
        return len(self.X)
    def __getitem__(self, idx):
        x = self.X[idx].astype(np.float32)
        y = self.Y[idx].astype(np.float32)
        return x, y

bostonds = SklearnDataset(X, y)
train_loader = DataLoader(bostonds, batch_size=32, shuffle=True, drop_last=True, )

class LinRegModel(pl.LightningModule):
    def __init__(self, input_dim: int):
        super().__init__()
        self.linear = nn.Linear(in_features=13, out_features=1, bias=True)

    def forward(self, x):
        y_hat = self.linear(x)
        return y_hat

    def training_step(self, batch, batch_idx):
        x, y = batch
        # flatten any input
        x = x.view(x.size(0), -1)
        y_hat = self(x)
        loss = F.mse_loss(y_hat, y, reduction="sum")
        return loss

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

trainer = pl.Trainer()
model = LinRegModel(input_dim=13)
trainer.fit(model, train_loader)

 

 

 

__init__ 함수 밑의 add_model_specific_args 함수는 ArgsBase의 함수와 마찬가지로 인자를 더해주는 함수인 것으로 보인다. 이제 그 밑에 적혀있는 두 함수에 대해 알아보겠다.

 

 

 

    def setup_steps(self, stage=None):
        # NOTE There is a problem that len(train_loader) does not work.
        # After updating to 1.5.2, NotImplementedError: `train_dataloader` · Discussion #10652 · PyTorchLightning/pytorch-lightning https://github.com/PyTorchLightning/pytorch-lightning/discussions/10652
        train_loader = self.trainer._data_connector._train_dataloader_source.dataloader()
        return len(train_loader)

    def configure_optimizers(self):
        # Prepare optimizer
        param_optimizer = list(self.model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(
                nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if any(
                nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters,
                          lr=self.hparams.lr, correct_bias=False)
        num_workers = self.hparams.num_workers

        data_len = self.setup_steps(self)
        logging.info(f'number of workers {num_workers}, data length {data_len}')
        num_train_steps = int(data_len / (self.hparams.batch_size * num_workers) * self.hparams.max_epochs)
        logging.info(f'num_train_steps : {num_train_steps}')
        num_warmup_steps = int(num_train_steps * self.hparams.warmup_ratio)
        logging.info(f'num_warmup_steps : {num_warmup_steps}')
        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=num_warmup_steps, num_training_steps=num_train_steps)
        lr_scheduler = {'scheduler': scheduler, 
                        'monitor': 'loss', 'interval': 'step',
                        'frequency': 1}
        return [optimizer], [lr_scheduler]

 

 

우선 setup_steps 에서는 인자로 받은 pl.Trainer 객체의 _data_connector_._train_dataloader_source.dataloaer() 함수를 사용하여 데이터를 불러오는 것으로 보인다. (그리고 데이터의 사이즈를 반환한다.)

 

 

 

configure_optimizers 함수에서는 우선 model.named_parameters() 함수를 사용하여 list 형태로 param_optimizer 변수를 받아준다. named_parameter 함수는 Module의 parameter에 대한 iterator을 반환하며 매개변수의 이름과 변수 자체를 모두 반환한다. 그 후 적혀있는 for 문은 조금 복잡하게 작성되어 있는데, 간단하게 글로 설명하자면 no_decay 리스트 안에 명시된 매개변수가 아닌 것들을 가져오는 것이다. (any 함수는 리스트 안에서 or 연산으로, all 함수는 리스트 안에서 and 연산으로 boolean(True/False) 값을 반환한다.) 

 

 

 

그 다음 만들어낸 두 개의 grouped_parameters 를 사용하여 AdamW를 사용한다. 이번에 AdamW에 대해서 처음 알아보게 되었는데, 우선 이에 대한 간단한 핵심정리 내용만 찾아서 읽어보았다. AdamW는 Decoupled Weigth Decay Regularization이라고도 부른다. Weight Decay(Weight들의 값이 증가하는 것을 제한함으로써 모델의 복잡도를 감소시켜 제한하는 기법)는 Loss Function에 L2 regularization을 추가하여 구현할 수 있으며 딥러닝 라이브러리가 optimization 함수에 동일한 방법으로 적용되어 있다고 한다. 하지만 Adam의 경우에는 파라미터마다 학습률을 다르게 적용하여 L2 regularization을 적용하면 성능이 하락하기에 이 문제를 해결하고자 weight decay를 분리하여 따로 구현한 것이라고 한다.

 

 

 

그 후에는 num_workers, data_len, num_train_steps 값을 지정해주고 있다. 그 다음에 num_warmup_steps 라는 변수를 만들어준다. Warmup Schedular은 딥러닝 모델에서 학습을 시작할 때 모든 파라미터가 보통 랜덤 value이기 때문에 최종 값에서 멀리 떨어져있는 문제를 완화시키기 위해 제안된 것인데, 초반에는 지정한 것보다 작은 learning rate를 사용하고 training 과정이 안정되면서 지정한 learning rate로 전환시키는 방법이다. 이 값은 총 train_step * warmup_ratio 값을 int 형으로 바꿔서 지정되어 있는 것을 확인할 수 있다.

 

 

 

그 다음 사용된 것이 transformers 라이브러리의 get_cosine_schedule_with_warmup 함수로, 안에 optimizer, num_warmup_steps, num_train_steps를 넣어준 것을 확인할 수 있다. 위 함수를 통해 만든 schedular과, monitor : loss, interval : step, frequency : 1로 지정하여 lr_schedular(Learning Rate Schedular)을 만들어주고, 최종적으로는 optimizer과 lr_schedular을 각각 리스트 안에 넣어서 반환하는 것을 확인할 수 있다.

 

 

 

class KoBARTConditionalGeneration(Base):
    def __init__(self, hparams, trainer=None, **kwargs):
        super(KoBARTConditionalGeneration, self).__init__(hparams, trainer, **kwargs)
        self.model = BartForConditionalGeneration.from_pretrained('gogamza/kobart-base-v1')
        self.model.train()
        self.bos_token = '<s>'
        self.eos_token = '</s>'
        
        self.tokenizer = PreTrainedTokenizerFast.from_pretrained('gogamza/kobart-base-v1')
        self.pad_token_id = self.tokenizer.pad_token_id

    def forward(self, inputs):

        attention_mask = inputs['input_ids'].ne(self.pad_token_id).float()
        decoder_attention_mask = inputs['decoder_input_ids'].ne(self.pad_token_id).float()
        
        return self.model(input_ids=inputs['input_ids'],
                          attention_mask=attention_mask,
                          decoder_input_ids=inputs['decoder_input_ids'],
                          decoder_attention_mask=decoder_attention_mask,
                          labels=inputs['labels'], return_dict=True)


    def training_step(self, batch, batch_idx):
        outs = self(batch)
        loss = outs.loss
        self.log('train_loss', loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        outs = self(batch)
        loss = outs['loss']
        return (loss)

    def validation_epoch_end(self, outputs):
        losses = []
        for loss in outputs:
            losses.append(loss)
        self.log('val_loss', torch.stack(losses).mean(), prog_bar=True)

 

 

 

드디어 실질적으로 학습을 시킬 모델의 Class이다. 이 Class는 방금 위에서 작성한 Base Class를 상속받아 사용하며, 실질적으로 LightningModule의 기능을 만들어줄 Class인 것이다. 

 

 

 

우선 init 함수를 보면 model.train() 이라고 적혀있는데, 이는 학습을시키는 함수가 아니다. nn.Module, pl.LightningModule 에는 Train에서와 Evaluation에서 수행하는 다른 작업을 Switch 할 수 있도록 train, eval 함수를 구현해두었다. Train step에서는 Evaluation step에서와 다르게 Dropout, Batch Normalization을 진행한다. 이러한 작업이 train, eval 단계에서 다르게 적용될 수 있도록 위 함수를 사용할 수 있는 것이다. 즉, train() 함수를 사용해주면 학습준비가 잘 되는 것이다.

 

 

 

그 후, forward 함수에서 attention_mask, decoder_attention_mask를 만들어주게 된다. Attention_mask들은 들어온 input_id의 벡터에 pad_token_id를 인자로 한 ne 함수를 사용하고(같지 않은 곳은 True, 같은 곳은 False) 값을 float()로 바꿔주어 1과 0으로만 이루어진 attention_mask 벡터를 만들어낸다. 

 

 

 

그 뒤의 training_step, validation_step, validation_epoch_end 는 batch를 받아서 loss 값을 계산해내는 과정을 거치고 있다.

 

 

 

이제 마지막 단계이다. 해당 파일이 실행명령을 받았을 때 실질적으로 실행되는 코드이다.

 

 

 

if __name__ == '__main__':
    parser = Base.add_model_specific_args(parser)
    parser = ArgsBase.add_model_specific_args(parser)
    parser = KobartSummaryModule.add_model_specific_args(parser)
    parser = pl.Trainer.add_argparse_args(parser)
    tokenizer = PreTrainedTokenizerFast.from_pretrained('gogamza/kobart-base-v1')
    args = parser.parse_args()
    logging.info(args)

    dm = KobartSummaryModule(args.train_file,
                        args.test_file,
                        tokenizer,
                        batch_size=args.batch_size,
                        max_len=args.max_len,
                        num_workers=args.num_workers)
    
    checkpoint_callback = pl.callbacks.ModelCheckpoint(monitor='val_loss',
                                                       dirpath=args.default_root_dir,
                                                       filename='model_chp/{epoch:02d}-{val_loss:.3f}',
                                                       verbose=True,
                                                       save_last=True,
                                                       mode='min',
                                                       save_top_k=3)

    tb_logger = pl_loggers.TensorBoardLogger(os.path.join(args.default_root_dir, 'tb_logs'))
    lr_logger = pl.callbacks.LearningRateMonitor()
    trainer = pl.Trainer.from_argparse_args(args, logger=tb_logger,
                                            callbacks=[checkpoint_callback, lr_logger])

    model = KoBARTConditionalGeneration(args, trainer)
    trainer.fit(model, dm)

 

 

 

지금까지 작성한 함수들을 사용하여 우선 인자들을 추가해준다(add_model_specific_args, Trainer.add_argparse_args). 그 후, dataset.py에서 작성한 KobartSummaryModule Class로 데이터를 불러오고, pl.callbacks.ModelCheckpoint로 모델 체크포인트를 지정해준다. Trainer 객체 또한 from_argparse_args로 파라미터들을 집어넣으며 생성해주고, 이 Trainer을 사용하여 모델을 만들어준 후, trainer.fit(model, data(dm)) 을 최종적으로 실행시키며 학습이 진행되는 것이다.

 

 

 

학습을 마치게되면 지정한 경로에 ckpt 파일을 저장하게 될 것이다. 이는 다시 사용하기 위해서 get_model_binary.py 파일을 실행시켜서 학습한 model binary 추출 작업을 진행한 이후 가능하다. (PyTorch-Lightning Binary => Huggingface Binary 형태로) 아래와 같은 코드를 사용한다.

 

 

 

import argparse
from train import KoBARTConditionalGeneration
from transformers.models.bart import BartForConditionalGeneration
import yaml

parser = argparse.ArgumentParser()
parser.add_argument("--hparams", default=None, type=str)
parser.add_argument("--model_binary", default=None, type=str)
parser.add_argument("--output_dir", default='kobart_summary', type=str)
args = parser.parse_args()

with open(args.hparams) as f:
    hparams = yaml.load(f)
    
inf = KoBARTConditionalGeneration.load_from_checkpoint(args.model_binary, hparams=hparams)

inf.model.save_pretrained(args.output_dir)

 

 

 

./logs/tb_logs/default/version_0/hparams.yaml 파일을 읽어온 후, load_from_checkpoint, save_pretrained 함수를 사용한다.

 

 

 

마지막으로 추론 관련 코드이다.

 

 

 

import torch
import streamlit as st
from kobart import get_kobart_tokenizer
from transformers.models.bart import BartForConditionalGeneration

@st.cache
def load_model():
    model = BartForConditionalGeneration.from_pretrained('./kobart_summary')
    # tokenizer = get_kobart_tokenizer()
    return model

model = load_model()
tokenizer = get_kobart_tokenizer()
st.title("KoBART 요약 Test")
text = st.text_area("뉴스 입력:")

st.markdown("## 뉴스 원문")
st.write(text)

if text:
    text = text.replace('\n', '')
    st.markdown("## KoBART 요약 결과")
    with st.spinner('processing..'):
        input_ids = tokenizer.encode(text)
        input_ids = torch.tensor(input_ids)
        input_ids = input_ids.unsqueeze(0)
        output = model.generate(input_ids, eos_token_id=1, max_length=512, num_beams=5)
        output = tokenizer.decode(output[0], skip_special_tokens=True)
    st.write(output)

 

 

 

위 코드를 보면 streamlit이라는 라이브러리가 사용되었는데, 이는 python으로 데이터분석을 위한 웹앱을 쉽게 만들어주는 라이브러리라고 한다. python 코드 몇 줄이면 동작하는 웹 서비스를 만들 수 있다고 한다!  load_model의 함수 위에 @st.cache 가 Decorator로 들어가있는데, 이는 local cache에 불러온 정보를 저장할 수 있게끔 해준다고 한다. 

 

 

 

그 후, 이전에 진행했던 것과 같은 방법으로, encode, unsqueeze(0)으로 1차원 벡터에서 2차원으로 만들어준 후, generate를 진행한다. 그 후, 결과값을 decode하여 text로 결과를 확인할 수 있는 것이다. 

 

 

 

여기까지, KoBART 소스코드들에 대한 간단한 분석 및 공부를 진행해보았다.

 

 

 

 

출처:

https://abluesnake.tistory.com/128

https://docs.streamlit.io/library/api-reference/performance/st.cache

https://wikidocs.net/52846

https://minimin2.tistory.com/184

https://rabo0313.tistory.com/entry/Pytorch-modeltrain-modeleval-%EC%9D%98%EB%AF%B8

https://better-tomorrow.tistory.com/entry/Learning-rate-Warmup

https://deep-learning-study.tistory.com/750

https://hiddenbeginner.github.io/deeplearning/paperreview/2019/12/29/paper_review_AdamW.html

https://minimin2.tistory.com/41

https://docs.python.org/ko/3/howto/logging.html

https://wotres.tistory.com/entry/def-%EB%92%A4%EC%97%90-None-%EB%98%90%EB%8A%94-str-%EC%9D%98%EB%AF%B8

https://wikidocs.net/157586

https://pytorch-lightning.readthedocs.io/en/stable/common/hyperparameters.html