[Code review] Transformers Trainer class 기능 정리

Transformers (HuggingFace)

2019 ~ 2020년. BERT 논문이 나온 이후 관련된 오픈 소스 라이브러리는 마치 춘추전국시대가 연상되었다. 유저들은 facebook, nvidia, google 등 많은 곳에서 오픈 소스를 사용했다. 하지만 결국 승리자는 HuggingFace 라고 생각한다. 본 포스팅에서는 2020년도부터 행복하게 사용해온 HuggingFace 의 Transformers 라이브러리 중 학습 Main 모듈인 Trainer Class 를 살펴보려 한다.

• Gihub Link
• Documents Link

출처: https://huggingface.co/

---

Transformers > Trainer 객체 지정

    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        model_init: Optional[Callable[[], PreTrainedModel]] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
    ):

학습 Main 모듈인 Trainer 객체를 만들기 위해서는 위와 같은 인자를 정의해야 한다.
필수로 정의해야 하는 인자는 다음과 같다.

• model : 학습하고자 하는 모델
• args : TrainingArguments 로 정의한 학습 조건
• train_dataset : 학습하고자 하는 데이터
• eval_dataset : 검증하고자 하는 데이터 (정의하지 않아도 되지만 정의 추천)
• tokenizer : 모델 학습에 사용되는 tokenizer
• compute_metrics : 검증 방식 (정의하지 않아도 되지만 정의 추천)

그 외 data_collator, model_init, callbacks, optimizers, preprocess_logits_for_metrics 와 같은 인자는 docs 를 참고하길 바란다. 이 중 필자가 자주 사용한 인자는 callback 이다. callback 관련된 설명은 다음에 포스팅할 예정이다.

만약 Multi-GPUs 사용하여 모델을 학습하고자 한다면 자동으로 DP 방식을 사용한 학습을 진행하게 된다. 소스 코드로 보면 device 의 수가 1보다 크면 DP 사용으로 되어 있다. 다만 Model Parallel 또는 DDP 등 학습 기능을 사용할 수 있기 때문에 device 의 수가 1보다 크면 DP 사용이 강제는 아니다. 이와 관련된 예외 처리 또한 소스 코드에 있으니 궁금하면 확인해봐도 좋을 것 같다.

효율적으로 모델을 학습하는 방법으로 DP, MP, Deep speed, mixed precision 등 다양한 기능을 제공한다. 자신이 사용하는 노드의 스펙을 확인하여 효율적인 학습을 하기 위해 설정하는 걸 추천한다.

다양한 학습 방법 및 기능 사용은 TrainingArguments 에서 정의할 수 있다.

---

Transformers > train 메소드

    def train(
        self,
        resume_from_checkpoint: Optional[Union[str, bool]] = None,
        trial: Union["optuna.Trial", Dict[str, Any]] = None,
        ignore_keys_for_eval: Optional[List[str]] = None,
        **kwargs,
    ):

Trainer 객체를 지정했다면 Trainer.train() 을 통해 모델을 학습시킬 수 있다. train 메소드는 학습 Main 메소드로서 Transformers 의 많은 함수 또는 객체를 활용하여 모델을 학습시킨다. 소스 코드로 보면 몇 줄 안 되는 것 같은데 자세히 살펴보면 굉장히 많은 기능을 끌어다 사용하고 있다는 걸 알 수 있다.

• Resume 기능
  • 재학습 기능이다. 모델 학습이 중간에 멈췄을 때 내가 원하는 시점의 Save point 로부터 재학습을 진행한다.
  • Save point 폴더를 확인하면 많은 파일들이 저장된다. 이 파일들은 Resume 기능을 활용하기 위해 Training Arguments, Learning rate scheduler 등 중간 학습 상태를 저장하고 활요하기 위해 저장된다.
  • 만약 resume_from_checkpoint = True 로 설정한다면 가장 마지막 Save point 로부터 모델을 불러와 재학습을 진행한다. 물론 directory 로 설정하면 내가 원하는 Save point 로부터 재학습을 할 수 있다.
• 그 외 기능은 본격적으로 학습을 진행하기 위해 환경을 설정하는 단계라고 생각하면 된다.

train 메소드는 결국 _inner_training_loop 메소드를 실행시키기 위한 준비 단계이다.

---

Transformers > _inner_training_loop 메소드

    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):

본격적으로 학습이 진행되는 메소드이다.

Transformers 에는 학습 진행 중 log 로 많은 정보를 쏴준다. _inner_training_loop 메소드에서는 log 로 쏴줄 수 있는 많은 정보를 저장하고 callback(⊃ self.control) 을 통해 유저에게 다양한 정보를 보여준다.
어느 정도 학습이 진행됐는지(epoch, step, global_step 등), 학습이 진행되는 조건은 어떠한지(total train batch size, total optimization steps 등), 현재는 어느 단계인지를 log 로 유저에게 제공한다.

gradient_accumulation_steps 을 통해 많은 train batch size 로 모델 학습을 할 때 GPU 메모리를 절약할 수 있다. 세계 석학들의 연구로 batch size 를 높게 설정하면 transformer 기반 모델의 성능이 향상되는 걸 입증했다.
batch size 를 대폭 늘려 학습할 수 있도록 돕는 gradient_accumulation_steps 은 _inner_training_loop 메소드에서 활용되고 있다.

본격적으로 학습이 진행되는 _inner_training_loop 메소드도 다른 많은 메소드를 활용한다.
대표적인 몇 개의 메소드를 소개하겠다.

• training_step : loss 계산부터 backward 까지 계산을 담당
• _maybe_log_save_evaluate : eval point 별 성능을 계산 (self.control 로 계산 시점 관리)
• _load_best_model : best model 을 비교 및 load

학습 중간중간 이때가 save point 야, eval point 야 등을 알려주는 기능은 callback_handler 를 통해 정보를 관리한다. 해당 기능은 callback 기능을 소개하는 포스팅에서 자세히 다루도록 하겠다.

---

Transformers > training_step 메소드

    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:

model 및 inputs 사용하여 loss 를 계산한다. 더 큰 batch_size 를 활용하기 위해 gradient_accumulation_steps 을 사용하여 loss 를 계산할 수 있다.

loss 계산 방법은 compute_loss 메소드 참고

---

Transformers > compute_loss 메소드

def compute_loss(self, model, inputs, return_outputs=False):

loss 계산하는 메소드로 loss 뿐만 아니라 logtis, hidden_states, attentions 와 같은 outputs 값도 추출할 수 있다.
단, Transformer 계열의 모델을 사용할 때 hidden_states, attentions 값을 뽑아낼 때는 OOM 을 조심하자.

---

Transformers > _maybe_log_save_evaluate 메소드

def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch, ignore_keys_for_eval):

log 를 쏴주거나 metrics 계산을 하거나 모델을 save 해준다.

log 의 경우, logging_strategy arg 를 통해 어느 시점에 log 를 쏴줄지 정할 수 있고
eval 의 경우, eval_strategy arg 를 통해 어느 시점에 성능을 계산할지 정할 수 있고
save 의 경우, save_strategy arg 를 통해 어느 시점에 모델을 저장할지 정할 수 있다.

어느 시점에 이 모든 처리를 할지는 self.control 즉, callback_handler 를 통해 이루어진다.
자꾸 쓰지만 callback 관련된 설명은 다음 포스팅에서 다루도록 하겠다.
사실.... callback 포스팅 빨리 쓰고 싶다. 너무 좋은 기능이라 생각하기 때문에...

---

Transformers > _load_best_model 메소드

def _load_best_model(self):

함수 이름처럼 best model 을 load 해준다.
물론 best model 은 모델의 성능 기준 값이 있어야 된다.
best model 의 dir 를 찾아 해당 dir 에서 모델의 weight 를 load 시켜주는 방식으로 best model load 가 진행된다.

---

Transformers > 그 외 학습 관련 메소드

그 외 Trainer.train 과 같이 학습과 관련된 메소드는 차암... 차~~~암 많다.
간단하게 몇 개만 소개하면

• _save_checkpoint
• _load_optimizer_and_scheduler
• log
• save_model
• _sorted_checkpoints
• 등등...

이 사용되고 있다.
궁금하면....
오배건~~ (ㅈㅅ...)

---

Comments

지금까지 HuggingFace 의 Transformers Trainer class 를 사용한 학습 소스 코드에 대해 알아보았다.
필자는 Transformers 패키지의 열혈 찬양을 하는 사람이다.
왜냐하면 어떠한 방법론, 모델이 들어와도 계속해서 업데이트하고 append 할 수 있는 구조로 되어있기 때문이다.
많은 연구자들이 이러한 설계를 했다고 생각한다.

혹시라도 모듈 개발을 해야 한다면 Transformers 패키지를 하나하나 뜯어보면서 어떻게 설계했는지 아이디어를 얻어 내가 개발하고자 하는 모듈을 계속 업데이트할 수 있도록 설계하는 걸 추천한다.

다음 포스팅에서는 Transformers 패키지의 callback 함수를 살펴볼 예정이다.
예고편으로 앞으로 살펴볼 Transformers 패키지의 기능을 미리 소개하겠다.

• class transformers.TrainerCallback
• class transformers.TrainerState
• class transformers.TrainerControl
• class transformers.Trainer > evaluate&predict
• class transformers.TrainingArguments
• class transformers.TrainerControl
• class transformers.BertModel

To Be Continued.....