注意
单击 此处 下载完整的示例代码
使用 XLM-RoBERTa 模型进行 SST-2 二元文本分类¶
作者: Parmeet Bhatia
概述¶
本教程演示了如何使用预训练的 XLM-RoBERTa (XLM-R) 模型在 SST-2 二元数据集上训练文本分类器。我们将展示如何使用 torchtext 库来
为 XLM-R 模型构建文本预处理管道
读取 SST-2 数据集并使用文本和标签转换对其进行转换
使用预训练的 XLM-R 编码器实例化分类模型
常用导入¶
import torch
import torch.nn as nn
DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
数据转换¶
像 XLM-R 这样的模型无法直接处理原始文本。训练这些模型的第一步是将输入文本转换为张量(数值)形式,以便模型能够对其进行处理并做出预测。处理文本的标准方法是
将文本进行分词
将分词转换为(整数)ID
添加任何特殊标记 ID
XLM-R 使用 sentencepiece 模型进行文本分词。下面,我们使用预训练的 sentencepiece 模型以及相应的词表,使用 torchtext 的转换构建文本预处理管道。使用 torchtext.transforms.Sequential()
对转换进行管道化,这类似于 torch.nn.Sequential()
,但支持 torchscriptable。请注意,转换支持批处理和非批处理文本输入,即可以传递单个句子或句子列表。
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.ac.cn/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.ac.cn/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]
脚本的总运行时间:(0 分钟 0.000 秒)