[Hands-On] BERT를 활용한 헤드 기반 텍스트 분류

Hugman Sangkeun Jung
25 min readApr 14, 2024

--

교육용 목적으로 작성된 코드입니다.

(영어버젼의 포스트는 링크에서 찾아볼 수 있습니다.)

분류 기법은 정말 다양합니다. 이 포스트에서는 그러한 분류 기술 중에 신경망, 특히 트랜스포머 기반의 분류 기술 중 “Head” 기반 분류 기법에 대한 실습을 진행해 보겠습니다. 이 글은 Head-based 분류 실습 시리즈 중 첫 번째 포스트로서, BERT를 활용한 텍스트 분류에 대해 살펴보겠습니다. 두 번째 포스트에서는 ViT에 기반한 이미지 분류를 살펴봅니다.

Head 기반 분류란 무엇인가요?

Head 기반 분류는 사전 훈련된 모델에 “헤드”(신경망 층 또는 층의 집합)를 추가하여 특정 작업을 수행하는 접근 방식입니다. 예를 들어, 텍스트의 주제를 분류하는 작업이 그것입니다. 여기서 “헤드”는 분류 작업에 대한 예측을 출력하기 위해 최적화되며, 충분한 데이터에 기반하여 네트워크가 해결하려는 문제 — 여기서는 분류 — 를 잘 해결할 수 있도록 풍부한 표현을 효과적으로 학습합니다.

헤드 기반 접근 방식을 활용함으로써 우리는 사전 훈련된 모델을 특정 작업에 적합하게 상대적으로 작은 데이터셋으로 미세 조정할 수 있고, 따라서 이 방법은 다양한 자연어 처리(NLP) 과제에 광범위하게 사용됩니다.

BERT와 Head 기반 분류

BERT 모델에서 [CLS] 토큰은 분류와 같은 시퀀스 수준 작업에서 중요한 역할을 합니다. 여기서는 [CLS] 토큰과 분류 헤드 간의 상호작용을 자세하게 단계별로 설명하겠습니다.

  1. 사전 훈련된 BERT 초기화
    풍부한 텍스트 코퍼스로부터 사전 훈련된 BERT 모델을 우선 준비합니다. 보통 huggingface 같은 도구를 이용해 다운로드 받습니다. 이 모델은 이미 범용 도메인에 대해 어느정도 문법, 맥락, 의미등을 포함해 언어를 이해할 수 있는 능력을 가지고 있습니다.
  2. [CLS] 토큰 추가
    각 입력 시퀀스에 [CLS] 토큰을 맨 앞에 추가합니다. 사실 표준 BERT의 경우 tokenizer를 구동시키는 순간 맨 앞에 [CLS] 라는 더미(dummy) 토큰이 추가되어 리턴됩니다. 이 [CLS] 토큰에 해당하는 위치에서의 Hidden state들은 BERT의 레이어를 통과하면서 입력 정보의 맥락을 집약하여 전체 시퀀스를 간략한 표현 — vector — 로 만들어 냅니다.
  3. 분류 헤드 설계
    BERT 의 최상단 레이어에서 [CLS] 토큰 위치에 해당하는 vector를 얻어내고, 이 vector를 분류 헤드에 연결합니다. 이 헤드는 종종 소프트맥스 레이어로 끝나는 간단한 신경망이며, [CLS] 토큰의 풍부한 표현을 특정 클래스 레이블로 매핑합니다.
  4. 미세 조정(Fine-tuning)
    분류 작업에 특화된 데이터셋에서 BERT와 분류 헤드를 결합한 모델을 미세 조정합니다. 이 단계는 특정 사용 사례에 잘 수행하도록 BERT의 파라미터와 분류 헤드 모두를 최적화합니다.
Head-based classification with BERT

구현 내용 미리보기

이번 실습에서는 다음과 같은 주요 단계를 거치게 됩니다.

  1. 데이터셋 준비 및 모델 초기화: AG News 데이터셋을 로드하여 GLUE 벤치마크의 일부로서 주제 분류를 위한 BERT 모델 훈련을 준비합니다. 이 과정에서 필요한 라이브러리와 도구를 준비하고, 사전 훈련된 BERT 모델을 초기화합니다.
  2. [CLS] 토큰과 분류 헤드 추가: 각 입력 시퀀스의 시작 부분에 [CLS] 토큰을 추가하고, BERT 모델의 출력에서 [CLS] 토큰에 대응하는 부분에 분류를 위한 헤드를 연결합니다. 이 분류 헤드는 네 가지 주제 카테고리를 반영하도록 설계됩니다.
  3. 미세 조정 및 평가: 사전 훈련된 BERT 모델과 추가된 분류 헤드를 AG News 데이터셋을 사용하여 미세 조정합니다. 훈련 과정은 배치별로 데이터를 모델에 공급하고, 에러를 역전파하여 모델의 가중치를 업데이트하는 사용자 정의 루프를 통해 진행됩니다. 검증 세트를 사용하여 모델 성능을 주기적으로 평가하고, 정확도, 정밀도, 재현율, F1 점수를 포함한 메트릭으로 최종 모델을 평가합니다.
  4. 결과 시각화 및 분석: 훈련된 모델의 성능을 혼동 행렬과 같은 시각적 도구를 사용하여 분석하고 시각화합니다. 이를 통해 모델이 각 주제에 대해 얼마나 잘 분류하는지, 어떤 주제에서 오류가 발생하는지 등을 파악할 수 있습니다.

환경 준비하기

데이터 로딩, 모델 구축, 훈련 및 평가에 필요한 라이브러리를 가져옵니다. 이 단계는 특정 작업에 필요한 도구와 라이브러리로 환경을 준비하는 중요한 과정입니다. 요리를 시작하기 전에 모든 재료를 모으는 것과 같습니다.

!pip install -qqq seaborn # for evaluation visualization
!pip install -qqq wandb # for logging
!pip install -qqq datasets # huggingface's lib.
!pip install -qqq transformers==4.39.2
!pip install -qqq accelerate==0.28.0
!pip install -qqq shortuuid

!pip install -U accelerate
!pip install tensorboard
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
from transformers import DataCollatorWithPadding
from datasets import load_dataset
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
import wandb
import torch
import random
import os

# Function to set the seed for reproducibility
def set_seed(seed_value=42):
"""Set seed for reproducibility."""
np.random.seed(seed_value)
torch.manual_seed(seed_value)
torch.cuda.manual_seed(seed_value)
torch.cuda.manual_seed_all(seed_value) # if you are using multi-GPU.
random.seed(seed_value)
os.environ['PYTHONHASHSEED'] = str(seed_value)

# The below two lines are for deterministic algorithm behavior in CUDA
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# Set the seed
set_seed()

데이터셋 로드하기

이 프로젝트에서 사용할 AG News 데이터셋은 Hugging Face의 datasets 라이브러리를 통해 로드합니다. 로드된 데이터셋은 훈련 데이터와 검증 데이터로 구성되어 있으며, 각 기사의 텍스트와 해당 기사의 주제를 나타내는 레이블로 구성됩니다. 데이터셋 로드 후, 데이터의 구조와 몇 가지 샘플을 확인하여 데이터셋이 올바르게 로드되었는지 확인합니다. 더 자세한 정보는 링크를 확인하세요.

dataset = load_dataset("ag_news")
print(dataset)
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 120000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 7600
})
})
from pprint import pprint # Using Python's pprint (pretty-print) library to prevent long horizontal output and make the data more readable

print(type(dataset)) # Data type
print(dataset) # Data structure and count
print("\n"*2+ "Train dataset:")
pprint(dataset["train"][1000]) # Print to check the content of train data

# Explanation of labels - 4 classes # https://huggingface.co/datasets/ag_news
# 1 class: World news
# 2 class: Sports news
# 3 class: Business news
# 4 class: Science/Technology news
class 'datasets.dataset_dict.DatasetDict'>
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 120000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 7600
})
})


Train dataset:
{'label': 3,
'text': 'European Union Extends Microsoft-Time Warner Review BRUSSELS, '
'Belgium (AP) -- European antitrust regulators said Monday they have '
'extended their review of a deal between Microsoft Corp. (MSFT) and '
'Time Warner Inc...'}

토크나이저 초기화 및 데이터 전처리

Tokenizer는 텍스트를 BERT 모델에 입력할 수 있는 토큰으로 변환합니다. 이 단계는 모델이 이해할 수 있는 형식으로 데이터를 준비합니다.

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def tokenize_function(examples):
return tokenizer(examples['text'], padding=True, truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

Data Collator

Data Collator 는 데이터를 특정 모델에 적합하게끔 전처리 해주는 역할을 자동으로 해줍니다. 기본적인 기능으로는 Batch 내의 입력길이에 맞춰 각 샘플의 문장길이를 Padding을 통해 맞춰줍니다.

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

사실 Data Collator는 huggingface에서 주로 사용하는 데이터 처리 패턴 중 하나입니다. 나중에 많이 쓰이게 될 형태이니, 한번 어떤 흐름으로 Huggingface내에서 데이터가 처리되는지 살펴보겠습니다.

Huggingface 에서의 Data 처리 흐름

거의 대부분의 경우 huggingface에서는 datasetmapping functiondata collatordata loader 의 흐름을 따라 진행합니다.

  1. dataset: AG News 데이터셋과 같은 데이터를 관리해주는 객체입니다.
  2. mapping function: 토큰화 및 전처리를 포함하여 원시 텍스트를 모델이 이해할 수 있는 형식으로 변환하거나 전처리하는 등의 기능을 구현합니다. 보통 개발자가 직접 데이터에 맞게 수정합니다.
  3. data collator: 가변 길이 시퀀스를 처리하는 데 필수적이며, 동적으로 배치를 패딩하여 균일성을 보장하고 모델의 원활한 입력을 촉진합니다. 보통 huggingface 에서 제공하는 모델형태(Seq2Seq model, TokenClassification, Causual_LM, …) 에 따라 그에 맞는 데이터 처리 기능이 이미 구현되어 배포됩니다. 필요시 직접 작성해도 됩니다.
  4. data loader: Batch, Shuffle, 모델에 대한 데이터 준비를 하는 최종 단계로, 훈련 또는 평가를 위한 최적화된 반복 가능한 객체를 생성합니다.

대게 이러한 흐름을 따라서 코딩을 수행하며, 이 패턴에 익숙해지면 다른 데이터셋이나 다른 모델을 활용할 때 높은 코드 가독성과 재활용성을 기대할 수 있습니다.

훈련하기

데이터가 준비되었으니, 이제 훈련을 진행해 봅시다. 아래의 단계를 통해 진행됩니다.

  1. 모델 정의: BERT를 로드하고 우리의 주제에 맞는 분류 레이어를 추가합니다.
  2. 훈련 인자: Epoch, Batch 크기, Learning rate과 같은 매개변수를 설정합니다.
  3. 훈련 루프: 데이터를 반복하며 파라미터를 조정하여 분류 정확도를 개선합니다.
  4. 평가: 모델의 성능을 훈련데이터가 아닌 따로 떼어놓은 데이터에 대해 테스트하여 얼마나 일반화 능력을 배워가는지 확인합니다.

모델 정의

여기에서 우리는 시퀀스 분류를 위한 BERT 모델을 불러와 초기화 하겠습니다. num_labels=4는 우리가 AG News 데이터셋의 topic수에 해당하는 것으로서, 즉 이렇게 불러온 BERT모델은 자동으로 Head를 붙여서 총 4개의 분류라벨을 출력할 수 있는 구조를 만들어 줍니다.

model = BertForSequenceClassification\
.from_pretrained('bert-base-uncased', num_labels=4)

훈련 인자 정의

훈련 인자는 에포크 수, 배치 크기, 학습률 조정 및 로그와 출력을 저장할 위치와 같은 훈련 과정에 대한 다양한 매개변수를 정의합니다.

training_args = TrainingArguments(
output_dir='./results-bert-topic-cls',
num_train_epochs=3,
per_device_train_batch_size=32,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',

evaluation_strategy='epoch', # Evaluate at the end of each epoch
logging_steps=10,
## ----
report_to="tensorboard",
)

평가용 함수 정의

보통 훈련루프를 정의하기 전에 훈련 도중에 자동으로 평가를 해주는 기능을 추가하기 위해 metric function 을 정의합니다.

def compute_metrics(p):
predictions, labels = p
predictions = np.argmax(predictions, axis=1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='weighted')
acc = accuracy_score(labels, predictions)
return {'accuracy': acc, 'f1': f1, 'precision': precision, 'recall': recall}

분류문제를 다루기 때문에, 여기서는 분류 테스크에서 많이 쓰이는 Precision, Recall, F1 Score, Accuracy를 활용하도록 하겠습니다. 각각의 의미는 여기서 자세하게 다루지는 않겠습니다. 더 깊은 이해를 원하시면 이 링크를 참고하세요

훈련 루프

# Select the first N samples from the tokenized training dataset
subset_train_dataset = tokenized_datasets['train'].select(range(6000)) # 1/2 data for time saving

trainer = Trainer(
model=model,
args=training_args,
train_dataset=subset_train_dataset,
eval_dataset=tokenized_datasets['test'],
data_collator=data_collator,
compute_metrics=compute_metrics,
)

전통적인 학습 과정에서는 for-loop를 통해 순전파와 역전파를 반복적으로 수행합니다. PyTorch에서 이 방식은 기본적이지만, Hugging Face의 Trainer 클래스를 사용하면 이 과정을 보다 쉽게 관리할 수 있습니다. Trainer는 학습, 평가, 예측을 위한 고수준 인터페이스를 제공하여 복잡한 학습 루프를 간단하게 만들어줍니다. 사용자는 TrainingArguments를 통해 학습 설정을 조정하고, Trainer의 메서드를 이용해 모델을 학습시키고 평가할 수 있습니다. 이는 학습 과정을 효율적으로 관리하고자 할 때 매우 유용한 도구입니다.

여기까지 준비가 되면 다음의 명령어를 통해 훈련을 수행합니다.

trainer.train()

빠른 학습을 위해 3 epoch 까지만 수행했습니다. 정밀하게 Hyperparameter를 건드리지 않았지만 약 90% 정도의 정확도를 보이게끔 학습이 되었음을 알 수 있습니다.

Model 저장하기

# Specify the directory where you want to save your model
output_dir = './bert-topic-cls'

# Save the model
model.save_pretrained(output_dir)
# Save the tokenizer
tokenizer.save_pretrained(output_dir)

방금 훈련한 모델과 tokenizer 를 저장합니다. 이렇게 저장해 놓으면 향후에 설명드릴 pipeline에서 바로 가져다 쓸 수 있습니다.

평가하기

우리가 훈련한 모델을 이제 훈련 해보겠습니다.

# Evaluate the model
results = trainer.evaluate()

trainer.evaluate() 를 수행하면 그 결과물로 results 는 아래와 같이 평가결과를 리턴합니다.

print( results )
{'eval_loss': 0.3104916214942932,
'eval_accuracy': 0.9064473684210527,
'eval_f1': 0.9062005426924208,
'eval_precision': 0.9085149708994765,
'eval_recall': 0.9064473684210527,
'eval_runtime': 244.8868,
'eval_samples_per_second': 31.035,
'eval_steps_per_second': 3.879,
'epoch': 3.0}

Confusion Matrix

분류 성능을 더 면밀하게 살펴보는 또 다른 방법으로 Confusion matrix가 있습니다. Confusion Matrix는 분류 문제에서 모델의 성능을 평가하는 데 매우 유용한 도구입니다. 이 행렬은 실제 클래스와 모델이 예측한 클래스를 기반으로, 예측의 정확성을 시각적으로 나타내줍니다. 행렬은 실제 클래스를 행으로, 예측된 클래스를 열로 구성하며, 각 셀의 값은 해당 클래스 조합에 속하는 샘플 수를 나타냅니다.

Confusion Matrix의 주 대각선(왼쪽 상단에서 오른쪽 하단으로 이어지는 대각선)은 모델이 올바르게 예측한 케이스, 즉 True Positives(TP)와 True Negatives(TN)를 나타냅니다. 반면, 비대각선 요소는 오류를 나타내는데, False Positives(FP)는 실제로는 부정적인 클래스를 긍정적으로 잘못 예측한 경우, False Negatives(FN)는 긍정적인 클래스를 부정적으로 잘못 예측한 경우를 의미합니다.

Confusion Matrix를 통해 정확도뿐만 아니라, 정밀도(Precision), 재현율(Recall), F1 점수 같은 다른 중요한 성능 지표들도 계산할 수 있습니다. 이러한 지표들은 모델의 성능을 더 면밀하게 평가하고, 특히 불균형 데이터셋에서 모델의 성능을 이해하는 데 도움을 줍니다.

우리의 results 와 정답 데이터를 활용하여 confusion matrix를 그려보죠.


# Predictions to get the confusion matrix
predictions = trainer.predict(tokenized_datasets['test'])
preds = np.argmax(predictions.predictions, axis=-1)
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix

label_map = {
'LABEL_0': 'World',
'LABEL_1': 'Sports',
'LABEL_2': 'Business',
'LABEL_3': 'Sci/Tech'
}

cm = confusion_matrix(predictions.label_ids, preds)

# label_map to labels
labels = [label_map[f'LABEL_{i}'] for i in range(len(label_map))]

# Confusion Matrix
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
plt.xlabel('Predicted labels')
plt.ylabel('True labels')
plt.title('Confusion Matrix with Label Names')
plt.show()
Confusion Matrix

좋은 Confusion Matrix는 대부분의 올바른 예측이 대각선에 집중되어 있어야 합니다. 이는 모델이 다양한 클래스를 정확하게 구분하고, 오류를 최소화했음을 의미합니다. 대각선의 값이 크고, 대각선 외의 값이 작을수록 모델의 성능이 좋다고 볼 수 있습니다.

예를 들어 위 그림에서는, 전반적으로 분류를 잘하고 있으나 전체적으로 보면, World topic 문장들을 Sci/Tech topic으로 혼동(confusion)하여 인식되고 있는것을 알 수 있죠. 만약 이 모델을 고도화 해야 한다고 하면, 저 부분에 있어서 더 신경을 써야 하겠죠.

예측하기

데이터셋 전체를 사용하는 대신, 개별 ‘문장’을 모델에 입력하여 예측하는 것도 데이터 준비 과정과 유사합니다. 이 과정에서는 먼저 토크나이저를 사용하여 문장을 토큰화합니다. 필요한 경우, 일정 수준의 전처리를 수행한 후, 문장을 모델에 순전파(feed-forward)하여 분류 레이어의 출력값을 얻어냅니다. 이 출력값을 분석하여 최종적으로 클래스 레이블을 도출합니다.

# Example sentence
sentence = "The stock market is reaching new heights."

# Tokenize the sentence
inputs = tokenizer(sentence, padding=True, truncation=True, max_length=512, return_tensors="pt")
import torch
# Move inputs to the same device as the model
inputs = {k: v.to(model.device) for k, v in inputs.items()}

# Make prediction
model.eval() # Set the model to evaluation mode
with torch.no_grad():
outputs = model(**inputs)
predictions = outputs.logits.argmax(-1).item() # Get the predicted class (index)

# Map the prediction index to the class name (if you have a label map)
simple_label_map = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
predicted_label = simple_label_map[predictions]

print(f"Sentence: '{sentence}'")
print(f"Predicted Label: '{predicted_label}'")
Sentence: 'The stock market is reaching new heights.'
Predicted Label: 'Business'

Hugging Face Pipeline 을 통해 예측해보기

직접 예측 코드를 작성해도 되지만, 일반적으로 예측 프로세스는 대부분 비슷하고 어느정도 패턴화 가능합니다. 이를 huggingface에서는 pipeline이라는 형태로 간소화 해두었습니다. 아래의 코드를 수행하면 바로 처리가 됩니다.

from transformers import pipeline

# Specify the path to your fine-tuned model or use a pre-trained model from the Hugging Face Model Hub
model_path = './bert-topic-cls' # Change this to your model's path or a Hugging Face model name

# Load the pipeline for text classification
classifier = pipeline("text-classification", model=model_path, tokenizer=model_path)

우리가 훈련한 모델과 tokenizer를 로딩하여 classifier를 준비하였습니다. 이제 이 classifier는 여러 문장에 대해 손쉽게 분류를 진행할 수 있습니다.

# Example sentences
sentences = [
"The stock market is reaching new heights.",
"The new sports car has been unveiled at the auto show.",
"The tech company announced its latest gadget yesterday."
]

# Make predictions
predictions = classifier(sentences)

# Print the predictions using the label map
for sentence, prediction in zip(sentences, predictions):
# Map the predicted label to the actual class name
class_name = label_map[prediction['label']]
print(f"Sentence: '{sentence}'")
print(f"Predicted Label: '{class_name}' with score {prediction['score']:.4f}\n")
Sentence: 'The stock market is reaching new heights.'
Predicted Label: 'Business' with score 0.9829

Sentence: 'The new sports car has been unveiled at the auto show.'
Predicted Label: 'World' with score 0.6967

Sentence: 'The tech company announced its latest gadget yesterday.'
Predicted Label: 'Sci/Tech' with score 0.9329

결론

이 글을 통해 우리는 head-based 분류 기술을 살펴보고, 이를 자연어 처리 분야의 특정 문제인 토픽 분류에 적용해 보았습니다. 특히, BERT 모델에서 [CLS ] 토큰에 해당하는 벡터를 활용하여 추가적인 분류 레이어를 부착하는 방법을 배웠습니다. 이 기술을 활용하여 AG News 데이터셋을 대상으로 모델을 훈련시키고 평가하였으며, 결과를 시각화하여 모델의 성능을 분석해 보았습니다.

이 튜토리얼 내용은 Colab에서 다운로드 받거나 실행해볼 수 있습니다.

--

--

Hugman Sangkeun Jung

Hugman Sangkeun Jung is a professor at Chungnam National University, with expertise in AI, machine learning, NLP, and medical decision support.