PyTorch로 텍스트 분류 구현

PyTorch로 텍스트 분류 구현

자연어 처리(NLP)에서 텍스트 분류는 중요한 작업 중 하나입니다. 예를 들어, 스팸 탐지, 감정 분석, 뉴스 분류 등 다양한 응용 분야에서 사용됩니다. 본 포스팅에서는 PyTorch를 사용하여 간단한 텍스트 분류 모델을 구현하는 방법을 설명합니다.

1. 텍스트 분류 개요

텍스트 분류는 주어진 문장을 특정 범주로 분류하는 작업입니다. 기본적인 접근 방식은 다음과 같습니다.

1. 텍스트 데이터를 숫자로 변환 (토큰화 및 임베딩)
2. 신경망 모델을 설계하여 학습
3. 테스트 데이터로 모델 평가

이번 예제에서는 IMDb 영화 리뷰 데이터를 이용하여 감정 분석을 수행하는 모델을 구축합니다.

2. 데이터 준비

먼저, torchtext 라이브러리를 사용하여 IMDb 데이터셋을 로드합니다.

import torch
import torch.nn as nn
import torch.optim as optim
import torchtext
from torchtext.datasets import IMDB
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import DataLoader

# 토크나이저 설정
tokenizer = get_tokenizer("basic_english")

# 데이터 로드
train_iter, test_iter = IMDB(split=("train", "test"))

3. 텍스트 전처리 및 벡터화

모델 학습을 위해 텍스트 데이터를 정수 인덱스로 변환해야 합니다.

# 어휘 사전 구축
def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)

vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>"])
vocab.set_default_index(vocab["<unk>"])

# 텍스트를 텐서로 변환하는 함수
def text_pipeline(text):
    return vocab(tokenizer(text))

def label_pipeline(label):
    return 1 if label == "pos" else 0

4. 데이터로더 구성

PyTorch의 DataLoader를 사용하여 배치 데이터를 생성합니다.

def collate_batch(batch):
    text_list, label_list = [], []
    for label, text in batch:
        text_list.append(torch.tensor(text_pipeline(text), dtype=torch.int64))
        label_list.append(torch.tensor(label_pipeline(label), dtype=torch.int64))
    return torch.nn.utils.rnn.pad_sequence(text_list, batch_first=True), torch.tensor(label_list)

batch_size = 16
train_dataloader = DataLoader(train_iter, batch_size=batch_size, collate_fn=collate_batch)

5. 모델 설계

간단한 RNN 기반의 텍스트 분류 모델을 설계합니다.

class TextClassificationModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim):
        super(TextClassificationModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, text):
        embedded = self.embedding(text)
        output, (hidden, _) = self.rnn(embedded)
        return self.fc(hidden[-1])

vocab_size = len(vocab)
embed_dim = 100
hidden_dim = 128
output_dim = 1

model = TextClassificationModel(vocab_size, embed_dim, hidden_dim, output_dim)

6. 모델 학습

손실 함수 및 옵티마이저를 설정하고 모델을 학습합니다.

criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def train_model(model, dataloader, criterion, optimizer, epochs=5):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for text, labels in dataloader:
            optimizer.zero_grad()
            predictions = model(text).squeeze(1)
            loss = criterion(predictions, labels.float())
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")

train_model(model, train_dataloader, criterion, optimizer)

7. 모델 평가

학습된 모델을 사용하여 테스트 데이터에서 정확도를 평가합니다.

def evaluate_model(model, dataloader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for text, labels in dataloader:
            outputs = model(text).squeeze(1)
            predictions = torch.round(torch.sigmoid(outputs))
            correct += (predictions == labels).sum().item()
            total += labels.size(0)
    print(f"Accuracy: {correct / total:.4f}")

evaluate_model(model, train_dataloader)

8. 결론

이번 포스팅에서는 PyTorch를 사용하여 텍스트 분류 모델을 구현하는 방법을 설명했습니다. IMDb 영화 리뷰 데이터를 이용하여 감정 분석 모델을 구축하였으며, RNN을 활용하여 텍스트를 분류하는 기법을 살펴보았습니다.