RoBERTa를 이용한 감정 분류

실행환경: Colab Pro (Python 3.7.15)

코랩은 90분 이상 동작이 없으면 런타임이 끊어질 수 있기 때문에 console을 열어 아래 스크립트를 실행한다. 

function connectRuntime(){
    button = document.querySelector("body > div.notebook-vertical > colab-status-bar").shadowRoot.querySelector("button");
    button.click();
    console.log(" --Connected");
}
setInterval(connectRuntime, 10 * 60 * 1000);

주기적으로 연결 버튼을 클릭해 코랩이 종료되지 않도록 한다. 

---

준비

!pip install transformers
!pip install datasets

import os
import re
import warnings
from google.colab import drive

import torch
import pandas as pd
from tqdm import tqdm
from datasets import Dataset
from transformers import (
    Trainer,
    TrainingArguments,
    RobertaTokenizerFast,
    RobertaForSequenceClassification,
    EarlyStoppingCallback
)
from sklearn.metrics import accuracy_score, f1_score
from sklearn.model_selection import train_test_split 

drive.mount("/content/drive")
warnings.filterwarnings("ignore")

(Colab 기준) 필요한 패키지는 위와 같다. 데이터를 불러오거나 모델을 저장하기 위해 Google Drive를 연결하는 작업도 해주었다. 

 

CWD = "/content/drive/MyDrive/DACON"

def join_path(*args):
    return os.path.join(CWD, *args)

DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
TRAIN_CSV = join_path("data", "train.csv")
TEST_CSV = join_path("data", "test.csv")
MODEL = "tae898/emoberta-large"
BATCH_SIZE = 32
EPOCHS = 20
MAX_LENGTH = 128

데이터 경로, 모델 이름과 같이 코드 전반에 사용되는 값들을 먼저 정의하였다. 

데이터는 '월간 데이콘 발화자의 감정인식'을 사용하였다. 

모델은 emoberta-large를 사용하였다. 해당 모델은 RoBERTa를 기반으로 감정 분류 발화 데이터(MELD, IEMOCAP)를 학습하였고, 예측 레이블의 개수가 사용하고자 하는 데이터와 동일했기 때문에 현재 풀고자 하는 감정 분류 문제에 적합할 것으로 생각했다. sibert, finbert 등 여러 모델을 사용해 보았지만 예측 레이블의 개수가 다르거나 감정 데이터로 사전 학습되지 않은 모델은 좋은 성능을 보이지 않았다.

• github: tae898/erc
• paper: https://arxiv.org/abs/2108.12009

EmoBERTa: Speaker-Aware Emotion Recognition in Conversation with RoBERTa

파라미터 설정은 'A Comparison of BERT...' 글을 참고하였으며, 이후 batch와 epoch을 변경해가며 실험해 보았다. 

 

TRAIN_ARGS = TrainingArguments(
    output_dir=join_path("emoberta"),
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    learning_rate=1e-6,
    warmup_steps=500,
    weight_decay=0.01,
    dataloader_num_workers=0,
    save_total_limit=1,
    load_best_model_at_end=True,
    evaluation_strategy="epoch",
    save_strategy="epoch"
)

그리고 모델의 파라미터를 하나의 객체로 묶어 전달한다. 이때 epoch, batch 등 파라미터를 설정할 수 있다.

evaluation_strategy를 설정하면 Training-Loss와 함께 Validation-Loss를 확인할 수 있다. 두 값의 비교를 통해 현재 모델이 over/under-fitting 되었는지 모니터링할 수 있다. (i.e. Training-Loss < Validation-Loss; over-fitting)

만약 Colab 환경에서 모델을 저장하며 학습한다면 save_total_limit 값을 1로 설정해 주어야 한다. 그렇지 않으면, 체크 포인트마다 학습된 모델을 각각 저장하기 때문에 디스크가 넘쳐 에러가 발생할 수 있다. 만약 디스크가 넘쳤다면 런타임을 해제한 후 다시 연결해 학습을 진행해야 한다. 

---

데이터 전처리

데이터셋은 영어 대화 데이터이며, ID(행 번호), Utterance(발화) 두 항목만 사용했다.

ID

Utterance

...

 

class LabelEncoder(object):
    def __init__(self):
        self._targets = [
            "neutral",
            "joy",
            "surprise",
            "anger",
            "sadness",
            "disgust",
            "fear",
        ]
        self.target_size = len(self._targets)

    def encode(self, labels):
        labels = [self._targets.index(lb) for lb in labels]
        return labels

    def decode(self, labels):
        labels = [self._targets[lb] for lb in labels]
        return labels

데이터셋의 레이블이 문자열 형식(i.e. "neutral")으로 되어 있기 때문에 데이터를 토큰화하고 모델에 학습시키기 위해 정수로 변환해 주어야 한다. LabelEncoder는 값을 받아 정수 레이블로 변환해 주는 객체이다. 사전 학습 모델의 장점을 살리기 위해 사전 학습 데이터에 사용된 레이블 인덱스를 유지하도록 하였다. 

 

train_csv = pd.read_csv(TRAIN_CSV)

label_encoder = LabelEncoder()
label_size = label_encoder.target_size
train_csv["Target"] = label_encoder.encode(train_csv["Target"])

train_csv = train_csv.loc[:, ["Utterance", "Target"]]
train_csv.rename(columns={"Target": "label"}, inplace=True)

먼저 CSV 파일을 불러와 DataFrame으로 만들었다. 

데이터셋의 레이블을 정수 레이블로 변환해 준다. 그리고 모델에 레이블 개수를 전달하기 위해 label_size에 개수를 할당했다. 

정수 레이블로 변환된 데이터셋은 발화 문장(Utterance)과 레이블(Target)만 추출한 후, "Target"를 "label"로 수정해 주었다. 

 

df_train, df_eval = train_test_split(train_csv, test_size=0.2, shuffle=True, random_state=42)

변환된 데이터를 랜덤하게 8:2 비율의 학습 데이터와 검증 데이터로 나누었다. 

---

모델 학습

def roberta_tokenize(data):
    return roberta_tokenizer(
        data["Utterance"],
        max_length=MAX_LENGTH,
        padding="max_length",
        truncation=True,
    )


train_set = Dataset.from_pandas(df_train.reset_index(drop=True))
eval_set = Dataset.from_pandas(df_eval.reset_index(drop=True))

train_set = train_set.map(roberta_tokenize, batched=True, batch_size=len(train_set))
eval_set = eval_set.map(roberta_tokenize, batched=True, batch_size=len(eval_set))

train_set.set_format("torch", columns=["input_ids", "attention_mask", "label"])
eval_set.set_format("torch", columns=["input_ids", "attention_mask", "label"])

모델에 학습하기 전 문장을 토큰화하는 전처리가 필요하다. 레이블링 된 DataFrame을 Dataset(datasets) 객체로 변환한다. 만약 학습 과정에서 torch를 활용한다면 Dataset(torch) 객체로 변환해 주어도 된다. 

데이터셋을 토큰화한다. batch_size는 데이터 길이와 동일하게 설정해 준다. batch 크기가 데이터 전체 길이와 동일하다는 것은 데이터가 하나씩 나누어져 있다는 의미로 학습 시에 사용하는 batch 크기와는 다르다. 

그리고 반드시 열 이름을 ["input_ids", "attention_mask", "label"]로 변환해 주는 작업이 필요하다. 만약 학습 과정에서 torch를 활용한다면 불필요한 부분이다. 

 

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="macro")
    acc = accuracy_score(labels, preds)
    return {"f1-macro": f1, "accuracy": acc}

모델에서 측정할 평가 지표를 함수의 형태로 넘겨준다. 위 예시는 f1-macro 점수와 accuracy를 계산하도록 작성하였다. 

 

model = RobertaForSequenceClassification.from_pretrained(MODEL, num_labels=label_size)

trainer = Trainer(
    model=model,
    args=TRAIN_ARGS,
    compute_metrics=compute_metrics,
    train_dataset=train_set,
    eval_dataset=eval_set,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=1)]
)

HuggingFace에서 제공하는 Roberta 객체를 사용하였고, 위에서 설정된 파라미터와 전처리된 데이터를 Trainer객체에 전달해 준다. callbacks는 모델의 평가 지표를 이용해 over-fitting된다고 판단될 때, Early-Stopping하여 성능이 가장 좋은 모델을 반환한다. 자세한 내용은 TrainingArguments에서 설정할 수 있다. 현재 코드는 Loss 값을 통해 Early-Stopping하도록 하였다. 

 

trainer.train()

학습을 진행한다. 

 

model_eval = trainer.evaluate()
print("Accuracy: {:.5f}".format(model_eval["eval_accuracy"]))
print("F1-macro: {:.5f}".format(model_eval["eval_f1-macro"]))

앞에서 작성한 모델 평가 함수를 통해 평가 지표를 얻을 수 있다.

 

결과: Label에 대해

랜덤하게 정수 레이블을 설정했을 때보다 사전 학습 모델과 동일하게 레이블을 설정했을 때 더 높은 성능을 보였다. f1-macro 값이 평균적으로 약 0.05 정도 높아졌다.

결과: shuffle에 대해

데이터가 연속적인 대화로 이루어져 있기 때문에 랜덤하게 섞지 않고 학습했을 때, 학습 효과가 좋지 않을까 의문이 들었다. 모델 성능 평가는 큰 차이가 없었지만 테스트 데이터에 대한 f1-macro 점수가 약간 낮아졌다.

결과: batch에 대해

8, 16, 32 중 32일 때, 가장 좋은 성능을 보였다. 

 

trainer.save_model(join_path("emoberta"))

모델 저장할 경로를 전달해 학습된 모델을 pytorch_model.bin과 config.json 파일로 저장한다. 

---

분류 (예측)

def roberta_tokenize(data):
    return roberta_tokenizer(
        data["Utterance"],
        max_length=MAX_LENGTH,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )

test_csv = pd.read_csv(join_path("data", "test.csv"))
df_test = test_csv.loc[:, "Utterance"].to_frame()

test_set = Dataset.from_pandas(df_test.reset_index(drop=True))
test_set = test_set.map(roberta_tokenize, batched=True, batch_size=len(test_set))
test_set.set_format("torch", columns=["input_ids", "attention_mask"])

데이터를 불러와 앞에서 진행한 방식과 동일하게 전처리한다. test_set은 "label"이 없는 상태이기 때문에  ["input_ids", "attention_mask"]만 가지게 된다. 

 

def predict(model, test_set):
    model.to(DEVICE)
    model.eval()
    
    test_predict = []
    for data in tqdm(test_set):
        input_id = data["input_ids"].unsqueeze(0).to(DEVICE)
        mask = data["attention_mask"].unsqueeze(0).to(DEVICE)
        output = model(input_id, mask)
        y_pred = output.logits
        test_predict += y_pred.argmax(1).detach().cpu().numpy().tolist()
    return test_predict

데이터를 학습된 모델에 전달해 분류를 진행한다. argmax를 통해 가장 큰 logits값의 인덱스를 예측 값으로 처리한다. 

 

preds = predict(model, test_set)
preds = label_encoder.decode(preds)

모델을 통해 분류된 레이블은 정수 레이블이기 때문에 앞에서 사용한 LabelEncoder를 통해 다시 문자열 형식으로 변환해 주는 작업을 한다. 

---

전체 코드: github/Denev6 - ipynb