[huggingface🤗] 3.Trainer사용해보기

Trainer

Trainer는 말 그대로 모델 학습을 담당하는 핵심 클래스입니다. 

 

Trainer는 데이터 로딩 전처리, 모델 생성과 최적화, 학습 루프 관리, 중간 결과 보고, 학습 결과 저장 등의 다양한 기능을 수행할 수 있습니다. 

 

클래스 초기화 

trainer = Trainer(
    model=model,  # 학습할 모델
    args=args,  # TrainingArguments
    train_dataset=train_dataset,  # 학습에 사용될 데이터셋
    eval_dataset=val_dataset,  # 평가에 사용될 데이터셋
    tokenizer=tokenizer,  # 토크나이저
    compute_metrics=compute_metrics,  # 평가 지표 계산 함수
    callbacks=[early_stopping_callback],  # 콜백 함수 리스트
)

일반적으로 위와 같이 trainer를 초기화 시켜 줍니다. 

• model(PreTrainedModel or torch.nn.Module) : 학습시킬 모델의 이름입니다. 
• args : TrainingArguments로 정의한 arguments입니다. 이 부분은 아래에서 설명하겠습니다.
• train_dataset : 학습 데이터 셋입니다. 
• eval_dataset : 평가 데이터 셋입니다.
• tokenizer : 토크나이저 입니다.
• compute_metrics : 계산해줄 평가 함수들을 적어주면 됩니다.
• callbacks : 콜백 함수의 리스트입니다. 학습 도중에 특정 시점에서 실행될 함수를 등록할 수 있습니다.

여기서 compute_metrics는 따로 메소드로 구현해줘야하는데요. 

이 함수는 EvalPrediction이라는 개체를 입력으로 받아야합니다. EvalPrediction은 아래의 두 값을 가지고 있습니다. 

• predictions : 모델의 예측 결과값을 포함하는 텐서
• label_ids : 실제 레이블 값을 포함하는 텐서

따라서 아래와 같이 지표 계산 코드를 넣어주시면 됩니다. 

from sklearn.metrics import accuracy_score

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return {"accuracy": accuracy_score(y_true=labels, y_pred=predictions)}

 

이 클래스를 사용하기 전에 TrainingArguments에 대해서 잠시 설명하겠습니다. 

 

TrainingArguments

args=TrainingArguments(
        output_dir='./results',  # 결과물 저장 경로
        num_train_epochs=3,  # 총 학습 에폭 수
        per_device_train_batch_size=16,  # GPU 하나당 배치 크기
        warmup_steps=500,  # Learning rate scheduler에서 warmup할 스텝 수
        weight_decay=0.01,  # 가중치 감소를 위한 하이퍼파라미터
        logging_dir='./logs',  # 학습 로그 저장 경로
        logging_steps=1000,  # 학습 로그 출력 간격
    )

TrainingArguments는 Trainer를 사용하기 위해 필요한 인자들을 정의해 주는 부분입니다. 

• output_dir : 결과물을 저장할 경로입니다.(폴더는 자동 생성 됩니다.)
• num_train_epochs : 총 학습의 에폭 수입니다.
• per_device_train_batch_size : GPU 하나당 할당할 배치의 크기입니다.
• warmup_stems : warmup할 스텝 수를 뜻합니다.
• weight_decay : 가중치 감소를 위한 파라니터입니다.
• logging_dir : 로그를 어디에 저장할지 경로를 설정해 줍니다.
• logging_steps : 로그를 얼마에 한번 출력할지 설정해줍니다.

자 이제 다시 Trainer로 돌아가보도록 하겠습니다. 

Trainer에서 쓸 수 있는 메소드 들을 살펴봅시다. 

▶ train()

위에서 정의한 모델 학습을 실행합니다. 

 

▶ evaludate()

위에서 정의한 모델의 평가를 실행합니다. 

 

▶ predict(test_dataset)

위에서 정의한 모델의 예측을 실행합니다. 결과로는 predictions, labels가 나옵니다.  

 

▶ save_model(model_path)

위에서 정의한 모델을 model_path에 저장합니다.

 

▶ get_train_dataloader()

학습 데이터의 로더를 반환합니다. 

 

▶ get_eval_dataloader()

검증 데이터의 로더를 반환합니다. 

 

▶ get_test_dataloader()

테스트 데이터의 로더를 반환합니다. 

 

▶ log_metrics(metrics)

지정된 메트릭을 로그에 출력하는 함수입니다. 입력으로 metric을 주면 됩니다. 

 

▶ get_model()

학습된 모델을 반환하는 함수입니다.

 

이외에도 정말 많은 메소드들이 있습니다. 

다른 함수를 보고 싶은 분들은 여기에서 확인해주세요!

 

전체 코드

from datasets import load_dataset
from transformers import BertForSequenceClassification, BertTokenizer
from transformers import TrainingArguments, Trainer
from sklearn.metrics import accuracy_score
import numpy as np

# 샘플 데이터셋 로드
dataset = load_dataset('glue', 'mrpc')
train_dataset = dataset['train']
eval_dataset = dataset['validation']

# bert-base-uncased 모델과 토크나이저를 로드 
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# bert 인풋에 맞춰주기 위한 전처리 함수
def preprocess_function(examples):
    return tokenizer(examples['sentence1'], examples['sentence2'], padding='max_length', truncation=True)
    
# metric을 구하는 함수
def compute_metrics(eval_pred):
    
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=1)
    
    return {"accuracy": accuracy_score(y_true=labels, y_pred=predictions)}

# 학습을 위한 argument 정의
training_args = TrainingArguments(
        output_dir='./results',  # 결과물 저장 경로
        num_train_epochs=3,  # 총 학습 에폭 수
        per_device_train_batch_size=16,  # GPU 하나당 배치 크기
        warmup_steps=500,  # Learning rate scheduler에서 warmup할 스텝 수
        weight_decay=0.01,  # 가중치 감소를 위한 하이퍼파라미터
        logging_dir='./logs',  # 학습 로그 저장 경로
        logging_steps=1000,  # 학습 로그 출력 간격
    )

# 학습을 위한 trainer 정의
trainer = Trainer(
    model=model, # 위에서 불러온 모델 
    args=training_args, # 위에서 정의한 trainingarguments
    tokenizer=tokenizer, # 위에서 정의한 토크나이저
    compute_metrics=compute_metrics, # 정의해둔 compute_metrics를 가져옴
    train_dataset=train_dataset.map(preprocess_function, batched=True), # 위에서 불러온 train data
    eval_dataset=eval_dataset.map(preprocess_function, batched=True) # 위에서 불러온 eval data
)

# 학습 진행
trainer.train()

# 평가 진행
results = trainer.evaluate()
print(results)