使用 XLM-RoBERTa 模型進行 SST-2 二元文字分類

作者：Parmeet Bhatia

概述

本教學示範如何使用預先訓練的 XLM-RoBERTa (XLM-R) 模型在 SST-2 二元資料集上訓練文字分類器。 我們將展示如何使用 torchtext 函式庫來

1. 為 XLM-R 模型建構文字預處理流程

2. 讀取 SST-2 資料集並使用文字和標籤轉換對其進行轉換

3. 使用預先訓練的 XLM-R 編碼器實例化分類模型

常用匯入

import torch
import torch.nn as nn

DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

資料轉換

像 XLM-R 這樣的模型無法直接處理原始文字。 訓練這些模型的第一步是將輸入文字轉換為張量（數值）形式，以便模型可以處理這些張量來進行預測。 處理文字的標準方法是

1. 將文字標記化

2. 將標記轉換為（整數）ID

3. 新增任何特殊標記 ID

XLM-R 使用 sentencepiece 模型進行文字標記化。 下面，我們使用預先訓練的 sentencepiece 模型以及相應的詞彙表，使用 torchtext 的轉換來建構文字預處理流程。 轉換使用 torchtext.transforms.Sequential() 進行流水線化，這類似於 torch.nn.Sequential()，但可以進行 torchscript 化。 請注意，轉換支援批次和非批次的文字輸入，也就是說，可以傳遞單個句子或句子列表。

import torchtext.transforms as T
from torch.hub import load_state_dict_from_url

padding_idx = 1
bos_idx = 0
eos_idx = 2
max_seq_len = 256
xlmr_vocab_path = r"https://download.pytorch.org/models/text/xlmr.vocab.pt"
xlmr_spm_model_path = r"https://download.pytorch.org/models/text/xlmr.sentencepiece.bpe.model"

text_transform = T.Sequential(
    T.SentencePieceTokenizer(xlmr_spm_model_path),
    T.VocabTransform(load_state_dict_from_url(xlmr_vocab_path)),
    T.Truncate(max_seq_len - 2),
    T.AddToken(token=bos_idx, begin=True),
    T.AddToken(token=eos_idx, begin=False),
)


from torch.utils.data import DataLoader

或者，我們也可以使用預先訓練模型附帶的轉換，它可以立即完成上述所有操作

text_transform = XLMR_BASE_ENCODER.transform()

資料集

torchtext 提供了幾個標準的 NLP 資料集。 如需完整列表，請參閱 https://pytorch.dev.org.tw/text/stable/datasets.html 的文件。 這些資料集是使用可組合的 torchdata 資料管道建構的，因此支援使用使用者定義的函數和轉換進行標準的流程控制和映射/轉換。 下面，我們將示範如何使用文字和標籤處理轉換來預處理 SST-2 資料集。

備註

使用資料管道目前仍有一些注意事項。 如果您希望擴展此範例以包含洗牌、多重處理或分散式學習，請參閱 本說明 以獲取更多說明。

from torchtext.datasets import SST2

batch_size = 16

train_datapipe = SST2(split="train")
dev_datapipe = SST2(split="dev")


# Transform the raw dataset using non-batched API (i.e apply transformation line by line)
def apply_transform(x):
    return text_transform(x[0]), x[1]


train_datapipe = train_datapipe.map(apply_transform)
train_datapipe = train_datapipe.batch(batch_size)
train_datapipe = train_datapipe.rows2columnar(["token_ids", "target"])
train_dataloader = DataLoader(train_datapipe, batch_size=None)

dev_datapipe = dev_datapipe.map(apply_transform)
dev_datapipe = dev_datapipe.batch(batch_size)
dev_datapipe = dev_datapipe.rows2columnar(["token_ids", "target"])
dev_dataloader = DataLoader(dev_datapipe, batch_size=None)

或者，我們也可以使用批次 API（即對整個批次應用轉換）

def batch_transform(x):
    return {"token_ids": text_transform(x["text"]), "target": x["label"]}


train_datapipe = train_datapipe.batch(batch_size).rows2columnar(["text", "label"])
train_datapipe = train_datapipe.map(lambda x: batch_transform)
dev_datapipe = dev_datapipe.batch(batch_size).rows2columnar(["text", "label"])
dev_datapipe = dev_datapipe.map(lambda x: batch_transform)

模型準備

torchtext 提供了 SOTA 預先訓練的模型，可用於在下游 NLP 任務上進行微調。 下面我們使用具有標準基本架構的預先訓練的 XLM-R 編碼器，並附加一個分類器頭，以在 SST-2 二元分類任務上對其進行微調。 我們將使用函式庫中的標準分類器頭，但使用者可以定義自己的適當任務頭並將其附加到預先訓練的編碼器。 如需可用預先訓練模型的其他詳細資訊，請參閱 https://pytorch.dev.org.tw/text/main/models.html 的文件

num_classes = 2
input_dim = 768

from torchtext.models import RobertaClassificationHead, XLMR_BASE_ENCODER

classifier_head = RobertaClassificationHead(num_classes=num_classes, input_dim=input_dim)
model = XLMR_BASE_ENCODER.get_model(head=classifier_head)
model.to(DEVICE)

訓練方法

現在讓我們定義標準的優化器和訓練標準，以及一些用於訓練和評估的輔助函數

import torchtext.functional as F
from torch.optim import AdamW

learning_rate = 1e-5
optim = AdamW(model.parameters(), lr=learning_rate)
criteria = nn.CrossEntropyLoss()


def train_step(input, target):
    output = model(input)
    loss = criteria(output, target)
    optim.zero_grad()
    loss.backward()
    optim.step()


def eval_step(input, target):
    output = model(input)
    loss = criteria(output, target).item()
    return float(loss), (output.argmax(1) == target).type(torch.float).sum().item()


def evaluate():
    model.eval()
    total_loss = 0
    correct_predictions = 0
    total_predictions = 0
    counter = 0
    with torch.no_grad():
        for batch in dev_dataloader:
            input = F.to_tensor(batch["token_ids"], padding_value=padding_idx).to(DEVICE)
            target = torch.tensor(batch["target"]).to(DEVICE)
            loss, predictions = eval_step(input, target)
            total_loss += loss
            correct_predictions += predictions
            total_predictions += len(target)
            counter += 1

    return total_loss / counter, correct_predictions / total_predictions

訓練

現在我們有了訓練分類模型的所有要素。 請注意，我們可以直接迭代我們的資料集物件，而無需使用 DataLoader。 我們的預處理資料集將產生資料批次，這要歸功於我們應用的批次資料管道。 對於分散式訓練，我們需要使用 DataLoader 來處理資料分片。

num_epochs = 1

for e in range(num_epochs):
    for batch in train_dataloader:
        input = F.to_tensor(batch["token_ids"], padding_value=padding_idx).to(DEVICE)
        target = torch.tensor(batch["target"]).to(DEVICE)
        train_step(input, target)

    loss, accuracy = evaluate()
    print("Epoch = [{}], loss = [{}], accuracy = [{}]".format(e, loss, accuracy))

輸出

100%|██████████|5.07M/5.07M [00:00<00:00, 40.8MB/s]
Downloading: "https://download.pytorch.org/models/text/xlmr.vocab.pt" to /root/.cache/torch/hub/checkpoints/xlmr.vocab.pt
100%|██████████|4.85M/4.85M [00:00<00:00, 16.8MB/s]
Downloading: "https://download.pytorch.org/models/text/xlmr.base.encoder.pt" to /root/.cache/torch/hub/checkpoints/xlmr.base.encoder.pt
100%|██████████|1.03G/1.03G [00:26<00:00, 47.1MB/s]
Epoch = [0], loss = [0.2629831412637776], accuracy = [0.9105504587155964]