[PyTorch] PyTorch 2.0 text classification 구현하기 with cuda 11.8 (1)

안넝

이번엔 저번 포스트에서 언급했던대로 2023년 4월 말 기준 가장 최근에 release된 PyTorch 버전 2.0과 이와 호환되는 torchtext 버전을 통해 조금 다른 방법으로 text classification 문제를 해결해보려고 한다.

(사실 이 예제가 2.0에서만 돌려볼 수 있는 게 아니라, 기존에 작성했던 예제가 2.0에서 호환이 되지 않으므로 "2.0에서도" 돌아가는 예제를 구현해 보는 것이 목표이다. 이후 PyTorch 2.0 이상에서만 지원하는 model.compile( ) 함수를 사용한 예제도 만들어볼 예정이다.)

우선 PyTorch 2.0과 이런저런 패키지들을 깔기 위해서는 아래 포스팅을 참고해주면 된다.

2023.04.27 - [공부/환경설정] - PyTorch 2.0 GPU 사용 설정

일단 내가 쓰고 있는 PC는 rtx3080, 워크스테이션은 rtx4090을 사용중이다.
혹시 똑같이 해보는 중에 호환성에 문제가 있다면 말해주길 바란다.

이전 포스팅에서 구버전의 pytorch와 torchtext.legacy 안에 있는 Field를 활용해 데이터를 만들었는데, torchtext 버전 0.11.0 이후부터는 해당 방법 사용이 불가능한 것으로 확인되어 겸사겸사 모델도 조금 바꿔서 다른 방식으로 해보기로 했다.

pytorch 2.0 버전 자체가 내가 훈련소에 들어갔을 때 나왔기 때문에 아직 예제가 많이 없어서, PyTorch 튜토리얼을 참고하여 구현해보고 있다.

해당 문제를 해결하면서 torchtext 라이브러리를 사용하여 어떻게 텍스트 분류를 위한 데이터셋을 만드는지를 살펴보도록 한다.
 1. Iterator를 통해 raw data에 접근
 2. Raw한 텍스트 문장들을 모델 학습에 사용할 수 있는 torch.Tensor로 변환하는 데이터 처리 pipeline 생성
 3. torch.utils.data.Dataloader를 사용하여 shuffle하고 반복(iterate)하기

우선 torchtext 데이터셋에 접근하기 전에, torchdata를 설치하면 된다. 다른 패키지를 설치하면서 같이 설치해주자.

conda install pytorch torchvision torchtext torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

1. Raw data iterator에 접근

torchtext 라이브러리는 가공되지 않은 텍스트 문장들을 생성하는(yield) 몇 가지의 raw dataset iterator를 제공한다.
이번 포스트에서 활용할 "AG_NEWS" dataset iterator는 lable, tuple의 형태로 가공되지 않은 데이터를 만든다.

import torch
from torchtext.datasets import AG_NEWS
train_iter = iter(AG_NEWS(split='train'))

print(next(train_iter))
print(next(train_iter))
print(next(train_iter))
print(next(train_iter))
print(next(train_iter))
print(next(train_iter))
print(next(train_iter))

Raw data 접근 결과

2. Data 처리 pipeline 준비

vocab, word vector, tokenizer를 포함하여 torchtext 라이브러리의 가장 기본적인 요소들을 검토한다.
이것은 raw한 text 문자열에 대한 기본적인 데이터 처리 빌딩블록이다.

vocab과 tokenizer를 사용한 기본적인 NLP 처리를 해보자. 첫 단계는 raw training dataset으로 어휘집(vocab)을 만드는 것이다. 여기서 token의 목록 또는 iterator를 받는 내장 함수인 build_vocab_from_iterator 를 사용한다. 이 함수를 통해 사용자는 어휘집에 추가할 special symbol 같은 것들을 전달할 수 있다.

################
# load_data.py #
################

from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator

tokenizer = get_tokenizer('basic_english')
train_iter = AG_NEWS(split='train')

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>"])

vocab block을 print 해보면 문자열이 아닌 정수로 변환되어져있다.

vocab(['here', 'is', 'an', 'example'])

위 예시의 print 결과

Text pipeline과 Label pipeline을 raw dataset iterator로부터 정수로 변환하여 받아와보자.

text_pipeline = lambda x: vocab(tokenizer(x))
label_pipeline = lambda x: int(x)-1

 

3. Data batch와 반복자 생성

PyTorch에서 제공하는 torch.utils.data.DataLoader를 활용하여 Data batch를 생성한다. 데이터로더에 대해서 알아보려면 이 링크를 확인하면 된다. getitem()과 len()을 활용해 구현한 맵 형태(map-style) 데이터셋으로 동작하며, map 처럼 index / key로 데이터를 얻어올 수 있다. 또한 shuffle을 False로 설정하면 iterable한 데이터셋처럼 동작한다.

모델로 data를 보내기 전, collate_fn 함수는 DataLoader로부터 생성된 샘플 배치로 동작한다. collate_fn 함수의 입력은 DataLoader에 batch size가 있는 batch 데이터이며, collate_fn 함수는 이를 미리 선언된 데이터 처리 pipeline에 따라서 처리한다. 해당 함수가 top-level에서 정의되어있으면 모든 곳에서 해당 함수를 사용할 수 있다.

아래의 코드에서 original data batch의 텍스트 데이터는 list에 담겨 nn.EmbeddingBag의 입력을 위한 하나의 tensor로 concatenate (합쳐지게) 된다. offset은 text tensor에서 개별 시퀀스 시작 index를 표현하기 위한 구분자 tensor이다. Label은 개별 텍스트 항목의 label을 저장하는 tensor이다.

from torch.utils.data import DataLoader
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def collate_batch(batch):
    label_list, text_list, offsets = [], [], [0]
    for (_label, _text) in batch:
         label_list.append(label_pipeline(_label))
         processed_text = torch.tensor(text_pipeline(_text), dtype=torch.int64)
         text_list.append(processed_text)
         offsets.append(processed_text.size(0))
    label_list = torch.tensor(label_list, dtype=torch.int64)
    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
    text_list = torch.cat(text_list)
    return label_list.to(device), text_list.to(device), offsets.to(device)

train_iter = AG_NEWS(split='train')
dataloader = DataLoader(train_iter, batch_size=8, shuffle=False, collate_fn=collate_batch)

 

4. Model 정의

모델은 nn.EmbeddingBag 레이어와 classification을 위한 linear 레이어로 구성된다. 기본 모드가 "평균"인 nn.EmbeddingBag은 임베딩들의 Bag들의 평균 값을 계산한다. 이때 Text 항목들은 각자 길이가 다를 수 있지만, 해당 모듈은 Text의 길이를 offset으로 저장하고 있으므로 padding이 필요하지는 않다. 또한 EmbeddingBag은 임베딩 평균을 즉시 계산하기 때문에 Tensor들의 시퀀스를 처리할 때 성능이나 메모리 효율성 면에서도 좋은 모습을 보여준다. 우선 기본 모델은 아래와 같다.

################
### model.py ###
################

from torch import nn

class TextClassificationModel(nn.Module):

    def __init__(self, vocab_size, embed_dim, num_class):
        super(TextClassificationModel, self).__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=False)
        self.fc = nn.Linear(embed_dim, num_class)
        self.init_weights()

    def init_weights(self):
        initrange = 0.5
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()

    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        return self.fc(embedded)

기본 모델은 위와 같고, 안에 Layer를 추가하여 활용할 수도 있다.

5. 인스턴스 생성

AG_NEWS dataset에는 4종류의 label이 존재하기 때문에 num_class = 4 이다.
( Label 1: World / Label 2: Sports / Label 3: Business / Label 4: Sci & Tec)

임베딩 차원이 64인 모델을 만들어준다. Voc_size의 크기는 vocab의 길이와 같다. num_class는 label의 개수와 같다.

train_iter = AG_NEWS(split='train')
num_class = len(set([label for (label, text) in train_iter]))
vocab_size = len(vocab)
emsize = 64
model = Conv1D(vocab_size, emsize, num_class).to(device)

 

6. 모델 Train & Evaluate 함수 정의

이제 main.py를 만들고, train / evaluate 함수를 정의해보자.

#################
#### main.py ####
#################

from torch.utils.data.dataset import random_split
from torchtext.data.functional import to_map_style_dataset
import time
from model import *


def train(dataloader):
    model.train()
    total_acc, total_count = 0, 0
    log_interval = 500
    start_time = time.time()
    for idx, (label, text, offsets) in enumerate(dataloader):
        optimizer.zero_grad()
        predicted_label = model(text, offsets)
        loss = criterion(predicted_label, label)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
        optimizer.step()
        total_acc += (predicted_label.argmax(1) == label).sum().item()
        total_count += label.size(0)
        if idx % log_interval == 0 and idx > 0:
            elapsed = time.time() - start_time
            print('| epoch {:3d} | {:5d}/{:5d} batches '
                  '| accuracy {:8.3f}'.format(epoch, idx, len(dataloader),
                                              total_acc / total_count))
            total_acc, total_count = 0, 0
            start_time = time.time()


def evaluate(dataloader):
    model.eval()
    total_acc, total_count = 0, 0

    with torch.no_grad():
        for idx, (label, text, offsets) in enumerate(dataloader):
            predicted_label = model(text, offsets)
            loss = criterion(predicted_label, label)
            total_acc += (predicted_label.argmax(1) == label).sum().item()
            total_count += label.size(0)
    return total_acc / total_count

 

7. Split train / test dataset

AG_NEWS 데이터셋에는 원래 evaluation용 데이터가 포함되어있지 않기 때문에, 학습 데이터를 가지고 train & test 데이터셋으로 분할 해야한다. 분할 비율은 0.95:0.05로 하겠다. 쉽게 랜덤으로 분할시키기 위해 PyTorch의 라이브러리 중 하나인 torch.utils.data.dataset.random_split( ) 함수를 사용한다.

Loss_function은 torch.nn.CrossEntropyLoss( ) 를 사용하며, 이는 nn.LogSoftmax( )와 nn.NLLLoss( )를 각 클래스에 대해 합쳐놓은 방식의 손실 함수이다. Optimizer는 SGD를 사용했으며, 확률적 경사 하강법을 뜻한다. 초기 LR은 큰 값인 5로 설정했는데, 매 epochs를 진행하면서 학습률을 StepLR을 통해 조절해 나간다.

# Hyperparameters
EPOCHS = 10 # epoch
LR = 5  # learning rate
BATCH_SIZE = 64 # batch size for training

loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.1)
total_accu = None
train_iter, test_iter = AG_NEWS()
train_dataset = to_map_style_dataset(train_iter)
test_dataset = to_map_style_dataset(test_iter)
num_train = int(len(train_dataset) * 0.95)
split_train_, split_valid_ = \
    random_split(train_dataset, [num_train, len(train_dataset) - num_train])

train_dataloader = DataLoader(split_train_, batch_size=BATCH_SIZE,
                              shuffle=True, collate_fn=collate_batch)
valid_dataloader = DataLoader(split_valid_, batch_size=BATCH_SIZE,
                              shuffle=True, collate_fn=collate_batch)
test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE,
                             shuffle=True, collate_fn=collate_batch)

for epoch in range(1, EPOCHS + 1):
    epoch_start_time = time.time()
    train(train_dataloader)
    accu_val = evaluate(valid_dataloader)
    if total_accu is not None and total_accu > accu_val:
      scheduler.step()
    else:
       total_accu = accu_val
    print('-' * 59)
    print('| end of epoch {:3d} | time: {:5.2f}s | '
          'valid accuracy {:8.3f} '.format(epoch,
                                           time.time() - epoch_start_time,
                                           accu_val))
    print('-' * 59)

output

 

8. test data를 통해 모델 평가

print('Checking the results of test dataset.')
accu_test = evaluate(test_dataloader)
print('test accuracy {:8.3f}'.format(accu_test))

output

9. 임의의 뉴스를 통해 결과 확인하기

학습된 모델을 활용해 골프 뉴스를 잘 인식할 수 있는지 테스트해보자.

ag_news_label = {1: "World",
                 2: "Sports",
                 3: "Business",
                 4: "Sci/Tec"}

def predict(text, text_pipeline):
    with torch.no_grad():
        text = torch.tensor(text_pipeline(text))
        output = model(text, torch.tensor([0]))
        return output.argmax(1).item() + 1

ex_text_str = "MEMPHIS, Tenn. – Four days ago, Jon Rahm was \
    enduring the season’s worst weather conditions on Sunday at The \
    Open on his way to a closing 75 at Royal Portrush, which \
    considering the wind and the rain was a respectable showing. \
    Thursday’s first round at the WGC-FedEx St. Jude Invitational \
    was another story. With temperatures in the mid-80s and hardly any \
    wind, the Spaniard was 13 strokes better in a flawless round. \
    Thanks to his best putting performance on the PGA Tour, Rahm \
    finished with an 8-under 62 for a three-stroke lead, which \
    was even more impressive considering he’d never played the \
    front nine at TPC Southwind."

model = model.to("cpu")

print("This is a %s news" %ag_news_label[predict(ex_text_str, text_pipeline)])

output

잘 잡는다.