KE-T5 Classification Example

[nannullna]

model/ke-t5 브랜치에 ke-t5.ipynb 파일이 있습니다!

해당 파일 참고하시면 될 듯합니다...

일단 필요한 것들을 로딩하고...

import numpy as np
import pandas as pd

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

from torch.utils.data import DataLoader

from transformers import T5Tokenizer, T5EncoderModel
from transformers import AdamW

from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

MODEL_NAME = 'KETI-AIR/ke-t5-large'
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)

저희 조 베이스라인에 기반해서 데이터셋, 데이터로더를 만듭니다.

from dataset.dataset import BaselineDataset
from dataset.preprocessor.preprocessors import BaselinePreprocessor

dataset = BaselineDataset(data_dir="/opt/ml/dataset/")
preprocessor = BaselinePreprocessor()

dataset.set_tokenizer(tokenizer)
dataset.set_preprocessor(preprocessor)
dataset.preprocess()

def collate_fn(samples):
    input_ids = [sample['input_ids'] for sample in samples]
    attention_mask = [sample['attention_mask'] for sample in samples]
    labels = [sample['label'] for sample in samples]

    input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
    attention_mask = torch.nn.utils.rnn.pad_sequence(attention_mask, batch_first=True, padding_value=0)
    
    return {'input_ids': input_ids.contiguous(),
            'attention_mask': attention_mask.contiguous(),
            'label': torch.stack(labels).contiguous()}

BATCH_SIZE = 8

dataloader = DataLoader(dataset, BATCH_SIZE, shuffle=True, collate_fn=collate_fn)

이후 GRU와 FC 레이어를 제외한 모두를 freeze합니다.

for idx, (name, param) in enumerate(model.named_parameters()):
    if name.startswith("gru") or name.startswith("fc"):
        param.requires_grad = True
    else:
        param.requires_grad = False

optimizer와 criterion을 지정해두고...

LEARNING_RATE = 1e-3

optimizer = AdamW(model.parameters(),
                  lr=LEARNING_RATE,
                  weight_decay=0.01,
                  correct_bias=False)

criterion = nn.CrossEntropyLoss()

훈련을 시킵니다!

NUM_EPOCHS = 10
LOG_EVERY  = 100
losses = []

model.train()
for epoch in range(NUM_EPOCHS):
    print("="*10, epoch, "="*10)
    partial_loss = []
    for idx, inputs in tqdm(enumerate(dataloader)):
        
        input_ids = inputs['input_ids'].to(device)
        attention_mask = inputs['attention_mask'].to(device)
        labels = inputs['label'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()
        partial_loss.append(loss.item())

        if (idx+1) % LOG_EVERY == 0:
            cumulated_loss = np.sum(partial_loss)
            partial_loss = []
            losses.append(cumulated_loss)
            print(losses[-1])

loss를 그려봅니다

import matplotlib.pyplot as plt

plt.plot(losses)
plt.show()

inference와 결과는 내일 중에 추가해보겠습니다!

[nannullna]

T5 모델 자체를 finetuning하는 시도도 해볼만하겠지만... 아직 엄두가 나지 않습니다! BATCH_SIZE = 8로 했을 때, GPU 메모리 20기가 정도를 사용하고, 점유율도 대략 80% 후반이 나옵니다! 물론... T5 레이어까지 fine tuning하면 훨씬 더 오래걸리겠지만요...