Hugging Face Trainer로 작은 모델 Fine-tune하기

1. 개요

이 글에서는 작은 모델(DistilBERT)을 transformers 라이브러리의 Trainer API로 fine-tune하여 주문 리뷰의 감정을 분류하는 방법을 배웁니다. 글을 다 읽고 나면 데이터셋을 준비하고, 토큰화하고, TrainingArguments를 구성한 후, 모델을 훈련하고, 평가하고, 추론용으로 저장하는 방법을 알게 될 것입니다.

이 글에서 사용할 예제: OrderReview – 고객(Customer)이 주문(Order)에 대해 남긴 리뷰이며, positive/negative 레이블이 붙어 있습니다.

2. 환경 준비

필요한 라이브러리를 설치합니다. 강력한 GPU가 없어도 됩니다. DistilBERT 같은 작은 모델은 작은 데이터셋에서 CPU로도 충분히 돌아갑니다. 다만 느릴 뿐입니다.

pip install transformers datasets evaluate accelerate torch scikit-learn

참고: 머신에 GPU가 있다면 해당 CUDA 버전의 torch를 따로 설치하여 훈련 속도를 크게 높일 수 있습니다. CPU의 경우, 데이터 행 수와 에포크를 제한하여 데모가 너무 오래 걸리지 않도록 하세요.

3. 데이터셋 준비

거대한 데이터셋을 다운로드하는 대신, 주문 리뷰를 모방한 작은 집합을 만들어 흐름을 쉽게 이해할 수 있도록 합니다. 실제로는 load_dataset("imdb")나 자신의 CSV 데이터로 바꿀 수 있습니다.

from datasets import Dataset

reviews = [
    {"text": "주문이 빠르게 배송되고, 포장이 신중하게 처리되었고, 매우 만족합니다", "label": 1},
    {"text": "제품이 결함이 있고, 배송이 느리며, 설명과 다릅니다", "label": 0},
    {"text": "품질이 좋고, 가격이 합리적이며, 다시 구매할 것입니다", "label": 1},
    {"text": "서비스가 형편없고, 고객 지원이 느립니다", "label": 0},
    {"text": "이 주문에 매우 만족하며, 5점입니다", "label": 1},
    {"text": "상품을 받았을 때 움푹 들어가 있었고, 실망했습니다", "label": 0},
]

dataset = Dataset.from_list(reviews)
dataset = dataset.train_test_split(test_size=0.33, seed=42)
print(dataset)

DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 4})
    test: Dataset({features: ['text', 'label'], num_rows: 2})
})

Dataset.from_list는 딕셔너리 리스트를 Hugging Face의 Dataset으로 변환합니다. train_test_split은 훈련/테스트 집합으로 나눕니다. 실제 데이터는 몇 천 줄 이상이어야 하고, 이것은 처리 흐름을 보여주기 위한 예제일 뿐입니다.

4. 데이터 토큰화

모델은 텍스트를 직접 이해하지 못하므로, 토크나이저를 사용하여 input_ids와 attention_mask로 변환해야 합니다.

from transformers import AutoTokenizer

checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=32)

tokenized_dataset = dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0].keys())

dict_keys(['text', 'label', 'input_ids', 'attention_mask'])

• truncation=True: max_length보다 긴 문장을 자릅니다.
• padding="max_length": 모든 문장이 같은 길이가 되도록 토큰을 추가하여 배치 병렬 처리를 쉽게 합니다.
• map(..., batched=True): 배치 단위로 토큰화 함수를 적용하므로 행 단위 처리보다 빠릅니다.

참고: 영어 DistilBERT 모델은 실제로 한국어 리뷰를 fine-tune하기에 적합하지 않습니다. 여기서는 기술을 시연하기 위해 사용했습니다. 한국어 데이터의 경우 vinai/phobert-base 같은 체크포인트를 선택하세요.

5. 모델 로드 및 레이블 정의

from transformers import AutoModelForSequenceClassification

label2id = {"negative": 0, "positive": 1}
id2label = {v: k for k, v in label2id.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)

num_labels=2는 미리 훈련된 인코더 위에 2개의 출력(긍정/부정)을 가진 분류 헤드를 만듭니다. id2label과 label2id는 모델 출력이 0/1 같은 숫자 대신 사람이 읽을 수 있는 레이블을 반환하도록 합니다.

6. TrainingArguments 구성

TrainingArguments는 모든 하이퍼파라미터와 훈련 루프 구성을 담습니다.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./order-review-model",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=1,
)

• eval_strategy="epoch": 각 스텝 대신 각 에포크 후에 평가합니다.
• load_best_model_at_end=True: 훈련 완료 후 Trainer가 metric_for_best_model에 따라 최고 성능 체크포인트를 자동으로 로드합니다.
• per_device_train_batch_size=4: 데모 데이터셋도 작으므로 배치 크기를 4로 설정했습니다. 실제 데이터는 VRAM에 따라 16 또는 32로 늘립니다.

7. 메트릭 계산 함수 정의

Trainer는 모델을 어떤 기준으로 평가할지 알기 위해 compute_metrics 함수를 필요로 합니다.

import numpy as np
import evaluate

accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels)
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

eval_pred는 Trainer가 각 평가 후 전달하는 (logits, labels) 튜플입니다. np.argmax는 가장 높은 확률(로짓)을 가진 레이블을 예측값으로 선택합니다.

8. Trainer 초기화 및 훈련

from transformers import Trainer, DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

Epoch   Training Loss   Validation Loss   Accuracy   F1
1       0.693100        0.671204          0.500000   0.500000
2       0.654200        0.598311          1.000000   1.000000
3       0.581900        0.520447          1.000000   1.000000

DataCollatorWithPadding은 배치별로 동적 패딩을 처리합니다(토큰화 시 고정 max_length로 패딩하는 것보다 효율적이지만, 여기서는 이미 패딩되어 있으므로 주로 배치 텐서를 생성합니다). trainer.train()은 forward → backward → optimizer step → 손실 로깅을 실행하며, eval_strategy에 따라 자동으로 평가합니다.

8.1. 훈련 과정 모니터링

더 자세한 로깅을 원한다면 TrainingArguments에 report_to="tensorboard"를 추가한 후 tensorboard --logdir ./order-review-model/runs를 실행하세요. 데모 같은 작은 데이터셋의 경우 콘솔 로그만으로도 손실이 점진적으로 감소하는지 확인할 수 있습니다.

9. 평가 및 모델 저장

metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./order-review-model/final")
tokenizer.save_pretrained("./order-review-model/final")

{'eval_loss': 0.5204, 'eval_accuracy': 1.0, 'eval_f1': 1.0, 'eval_runtime': 0.12, ...}

trainer.evaluate()는 eval_dataset에서 다시 실행하고 메트릭 딕셔너리를 반환합니다. save_model은 구성과 가중치를 모두 저장합니다. 나중에 로드할 때 어휘 불일치를 방지하기 위해 토크나이저도 함께 save_pretrained하세요.

10. Fine-tune된 모델로 추론

pipeline을 사용하여 훈련된 모델을 빠르게 테스트합니다.

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="./order-review-model/final",
    tokenizer="./order-review-model/final",
)

result = classifier("배송이 정시에 도착했고, 제품이 기대 이상으로 아름답습니다")
print(result)

[{'label': 'positive', 'score': 0.9123}]

pipeline은 토큰화, forward 처리, softmax를 자동으로 처리하여 신뢰도를 나타내는 label과 score를 반환합니다. 이것은 모델을 PaymentService 같은 작은 서비스에 고객 피드백을 자동으로 분류하기 위해 통합하는 가장 빠른 방법입니다.

11. 모델 동작 검사 테스트 작성

데모 데이터셋으로 pytest를 사용하여 파이프라인이 올바르게 작동하고 체크포인트를 변경할 때 성능이 저하되지 않도록 하는 간단한 테스트를 작성할 수 있습니다.

import pytest
from transformers import pipeline

@pytest.fixture(scope="module")
def classifier():
    return pipeline(
        "text-classification",
        model="./order-review-model/final",
        tokenizer="./order-review-model/final",
    )

def test_positive_review_is_classified_correctly(classifier):
    result = classifier("제품이 훌륭하고, 배송이 빠릅니다")[0]

    assert result["label"] == "positive"
    assert result["score"] > 0.5

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("제품이 결함이 있고, 배송이 늦으며, 매우 실망했습니다")[0]

    assert result["label"] == "negative"
    assert result["score"] > 0.5

scope="module"은 파이프라인이 파일의 모든 테스트에 대해 한 번만 모델을 로드하도록 하여 각 테스트 케이스마다 로드하는 시간을 절약합니다. 이 테스트는 상위 수준의 동작만 확인하며(레이블이 올바른 방향인지), 모델 품질 평가에 사용하면 안 됩니다. 모델 품질 평가는 위의 평가 단계에 속합니다.

12. 일반적인 오류

• padding/truncation 누락: max_length보다 긴 문장은 배치 내 텐서 크기가 일치하지 않아 런타임 오류를 일으킵니다.
• num_labels 혼동: 데이터셋에 3개 클래스가 있지만 num_labels=2로 선언하면 Trainer는 여전히 실행되지만 분류 헤드가 출력을 충분히 제공하지 못해 결과가 완전히 잘못됩니다.
• 데이터셋이 너무 작으면서 실제 데이터에서 높은 F1을 기대하기: 이 글의 예제는 흐름을 설명하기 위한 것이며, 훈련 데이터 4줄은 모델이 일반화하기에 충분하지 않습니다. 여기서의 정확도=1.0을 기준으로 삼지 마세요.
• seed를 설정하지 않기: 훈련 결과는 실행할 때마다 달라져 하이퍼파라미터를 튜닝할 때 비교하기 어렵습니다. 훈련 전에 transformers에서 set_seed(42)를 사용하세요.
• 약한 머신에 큰 모델(base/large) 사용: distilbert/tiny 버전으로 시작하여 개발 루프를 빠르게 진행한 후, 파이프라인을 확정했을 때만 더 큰 모델로 전환하세요.

Trainer를 사용하지 말아야 할 때: 훈련 루프를 완전히 사용자 정의해야 하는 경우(예: 복잡한 다중 작업 손실, 표현하기 어려운 사용자 정의 스케줄러)가 있다면, 순수 PyTorch 루프를 작성하거나 accelerate를 직접 사용하는 것이 Trainer에 끼워 맞추는 것보다 더 유연합니다.

13. 결론

Trainer API는 훈련/평가 루프를 직접 작성할 필요 없이 작은 모델을 fine-tune할 수 있게 해줍니다. 데이터셋, 토크나이저, 모델, TrainingArguments만 준비하면 됩니다. 주문 리뷰 분류 같은 간단한 분류 문제의 경우 표준 흐름은 토큰화 → TrainingArguments 구성 → compute_metrics 정의 → Trainer 초기화 → train() → evaluate() → save_model()입니다. 실제 데이터가 더 복잡해지더라도 데이터셋과 체크포인트만 바꾸면 되고 나머지 파이프라인은 거의 그대로 유지됩니다.