Huggingface의 Transformer 라이브러리에서 pre-trained된 모델 활용하기

Transfer Learning은 많은 양의 데이터로 사전 훈련된 모델을 가져와 새로운 태스크를 학습시키는 기술이다. Hugging Face의 Transformers 라이브러리는 다양한 사전 훈련된 모델을 제공하며, 이러한 모델을 fine tuning하여 새로운 태스크를 수행할 수 있다.

이번 포스팅에서는 Hugging Face의 Transformers 라이브러리를 사용하여 Transfer Learning을 수행하는 방법에 대해 알아보자.

---

1. Downstream Task에 알맞게 pre-trained된 모델 수정

from_pretrained() 메서드를 사용하여 사전 훈련된 모델과 토크나이저를 불러올 수 있다. Downstream Task를 수행하기 위해서는 pre-trained 모델의 출력 차원을 조정해야 한다. 아래 코드는 BERT 모델을 불러와서 문장의 이진 분류를 수행하는 경우 사용된다.

from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 모델 불러오기
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-cased', num_labels=2) # 이진 분류 task를 수행하기 위해 모델의 출력 차원을 2차원으로 수정

Downstream Task에서 예측해야 할 클래스 수가 2개이므로 num_labels의 값을 2로 지정해 모델의 출력 차원을 2차원으로 수정한다.

2. 데이터셋 로딩 및 전처리

Pre-trained 모델을 활용하여 Downstream Task를 학습 및 추론하기 위해서는 입력 데이터를 모델이 이해할 수 있는 형태로 변환하는 전처리를 수행해야 한다.

from transformers import TextClassificationDataset

# 입력 데이터 토큰화 및 인코딩
train_text_dataset = TextClassificationDataset(tokenizer=tokenizer, file_path='train.csv', max_length=512)
val_text_dataset = TextClassificationDataset(tokenizer=tokenizer, file_path='val.csv', max_length=512)

3. Fine Tuning을 위한 Trainer 및 TrainingArguments 설정

fine tuning을 위해 Trainer 클래스와 TrainingArguments를 사용한다. Trainer 클래스는 fine tuning을 실행하고 모델 학습을 위한 인터페이스를 제공하고, TrainingArguments 클래스를 사용하여 Trainer에 필요한 인자를 설정할 수 있다.
Trainer와 TrainingArguments에 대한 자세한 내용은 아래 글에서 확인할 수 있다.

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

from transformers import Trainer, TrainingArguments

# 학습 설정을 TrainingArguments 클래스를 통해 정의
training_args = TrainingArguments(
    output_dir='./output',          # 학습 결과물을 저장할 디렉토리
    num_train_epochs=3,            # 학습 에폭 수
    per_device_train_batch_size=8, # 각 디바이스당 학습 배치 크기
    per_device_eval_batch_size=8, # 각 디바이스당 검증 배치 크기
    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_text_dataset,   # 학습 데이터셋 (torch.utils.data.Dataset)
    eval_dataset=val_text_dataset,     # 검증 데이터셋 (torch.utils.data.Dataset)
    compute_metrics=compute_pearson_correlation # 성능 평가 메소드
)

4. Trainer를 사용하여 Fine Tuning 실행

Trainer 객체를 생성하고 train() 메서드를 호출하여 fine tuning을 실행한다.

# 모델 학습
trainer.train()

5. Fine Tuned 모델 저장 및 추론

save_pretrained() 메서드를 사용하여 fine tuned 모델을 저장할 수 있다. 또한 evaluate() 메소드를 사용해서 추론을 진행할 수 있다.

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

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