AI/기타

HuggingFace Trainer로 모델 학습 및 평가하기

sangwonYoon 2023. 4. 18. 15:13

HuggingFace Trainer는 PyTorch와 같은 딥러닝 프레임워크에서 모델을 학습하기 위한 편리한 인터페이스를 제공하는 파이썬 라이브러리이다. Trainer 클래스 TrainingArguments 클래스는 Trainer를 사용하여 모델을 학습시키기 위한 설정과 제어를 담당하는 중요한 클래스들이다. 이 두 클래스를 통해 모델 학습의 다양한 설정을 조정하고, 학습 과정을 관리할 수 있다.


Trainer 클래스

Trainer 클래스는 모델의 학습을 관리하고 제어하는 클래스로, transformers 라이브러리에서 제공되는 핵심 클래스 중 하나이다. Trainer 클래스를 사용하면 모델 학습에 필요한 다양한 기능들을 활용할 수 있다.

주요 속성

  • model: 학습할 모델 객체. PyTorch와 같은 딥러닝 프레임워크의 모델 객체를 할당한다.
  • args: 학습 설정을 담고 있는 TrainingArguments 객체. 학습에 관련된 설정들을 조정하고 제어하는데 사용된다.
  • data_collator: 데이터를 배치(batch)로 변환하는데 사용되는 데이터 콜레이터(data collator) 객체. 데이터 로딩 및 배치 처리와 관련된 설정들을 제어하는데 사용된다.
  • train_dataset: 학습에 사용되는 데이터셋 객체. torch.utils.data.Dataset 클래스 객체이다.
  • eval_dataset: 평가에 사용되는 데이터셋 객체. torch.utils.data.Dataset 클래스 객체이다.
  • compute_metrics: 모델의 성능 평가를 위해 사용되는 메소드 또는 함수. 평가 지표를 계산하는 방식을 정의한다.
  • tokenizer: 모델의 토크나이저(tokenizer) 객체. 텍스트 데이터를 모델이 이해할 수 있는 형태로 변환하는데 사용된다.

주요 기능 및 메소드

  • train(): 모델을 학습시키는 메소드로, 학습 데이터셋과 학습 설정을 지정하여 모델을 학습시킬 수 있다.
  • evaluate(): 검증 데이터셋을 평가하는 메소드로, 모델의 성능을 평가할 수 있다.
  • predict(): 테스트 데이터셋에 대한 예측을 수행하는 메소드로, 모델의 예측 결과를 얻을 수 있다.
  • save_model(): 학습된 모델을 저장하는 메소드. 저장된 모델은 나중에 불러와서 추론(inference)이나 fine-tuning 등에 사용될 수 있다.

TrainingArguments 클래스

TrainingArguments 클래스는 Trainer 클래스의 학습 설정을 관리하는 클래스로, 모델 학습에 필요한 다양한 설정을 조정할 수 있다. Trainer 클래스의 train() 메소드를 호출할 때 TrainingArguments 클래스를 인자로 전달하여 학습 설정을 지정할 수 있다.

 

주요 속성

  • output_dir: 학습 결과물을 저장할 디렉토리를 지정한다.
  • num_train_epochs: 학습 에폭 수를 지정한다.
  • per_device_train_batch_size: 각 디바이스당 학습 배치 크기를 지정한다.
  • per_device_eval_batch_size: 각 디바이스당 검증 배치 크기를 지정한다.
  • save_steps: 일정한 스텝마다 모델 가중치를 저장하는 주기를 지정한다.
  • logging_dir, logging_steps: 로깅 관련 설정을 지정한다.

 

Trainer 클래스와 TrainingArguments 클래스는 함께 사용하여 모델의 학습 설정을 세부적으로 조정하고, 학습, 검증, 테스트 과정을 제어할 수 있다. 예를 들어, 다양한 학습 하이퍼파라미터들을 조정하고, 학습 과정에서의 로깅 및 모델 저장 설정을 관리하며, 학습 데이터와 검증 데이터의 배치 크기를 조정하고, 학습 에폭 수를 설정할 수 있다.

 

Trainer 클래스와 TrainingArguments 클래스를 사용하여 모델 학습을 관리하는 예시 코드는 아래와 같다.

from transformers import Trainer, TrainingArguments
from scipy.stats import pearsonr

# 성능 평가를 위해 피어슨 상관계수를 계산하는 함수
def compute_pearson_correlation(pred):
    preds = pred.predictions.flatten()
    labels = pred.label_ids.flatten()
    return {"pearson_correlation": pearsonr(preds, labels)[0]}

# 학습 설정을 TrainingArguments 클래스를 통해 정의
training_args = TrainingArguments(
    output_dir='./output',          # 학습 결과물을 저장할 디렉토리
    num_train_epochs=3,            # 학습 에폭 수
    per_device_train_batch_size=8, # 각 디바이스당 학습 배치 크기
    per_device_eval_batch_size=16, # 각 디바이스당 검증 배치 크기
    warmup_steps=1000,             # learning rate scheduler의 warmup step 수
    save_steps=500,                # 일정한 스텝마다 모델 가중치를 저장하는 주기
    logging_dir='./logs',          # 로깅 결과를 저장할 디렉토리
    logging_steps=100,             # 일정한 스텝마다 로깅을 수행하는 주기
)

# Trainer 클래스를 생성하고 학습 데이터와 학습 설정을 전달
trainer = Trainer(
    model=model,                   # 학습할 모델
    args=training_args,            # 학습 설정
    train_dataset=train_dataset,   # 학습 데이터셋 (torch.utils.data.Dataset)
    eval_dataset=eval_dataset,     # 검증 데이터셋 (torch.utils.data.Dataset)
    compute_metrics=compute_pearson_correlation # 성능 평가 메소드
)

# 모델 학습
trainer.train()

# 모델 저장
trainer.save_model('저장할 디렉토리 경로')

# 모델 평가
metrics = trainer.evaluate(eval_dataset=eval_dataset)

위와 같이 Trainer 클래스와 TrainingArguments 클래스를 사용하면, 모델의 학습을 세부적으로 제어하고, 학습 설정을 조정하여 최적의 학습 결과를 얻을 수 있다.