# 使用 XLM-RoBERTa 模型进行 SST-2 二元文本分类

> 原文：https://pytorch.org/text/stable/tutorials/sst2_classification_non_distributed.html

注意

点击这里下载完整的示例代码

作者：Parmeet Bhatia

## 概述

本教程演示了如何使用预训练的 XLM-RoBERTa（XLM-R）模型在 SST-2 二元数据集上训练文本分类器。我们将展示如何使用 torchtext 库来：

1.  为 XLM-R 模型构建文本预处理管道

1.  读取 SST-2 数据集并使用文本和标签转换进行转换

1.  使用预训练的 XLM-R 编码器实例化分类模型

## 常见导入

```py
import torch
import torch.nn as nn

DEVICE = [torch.device](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")("cuda") if [torch.cuda.is_available](https://pytorch.org/docs/stable/generated/torch.cuda.is_available.html#torch.cuda.is_available "torch.cuda.is_available")() else [torch.device](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device "torch.device")("cpu") 
```

## 数据转换

像 XLM-R 这样的模型不能直接处理原始文本。训练这些模型的第一步是将输入文本转换为张量（数值）形式，以便模型可以对其进行处理并进行预测。处理文本的标准方法是：

1.  对文本进行标记化

1.  将标记转换为（整数）ID

1.  添加任何特殊的标记 ID

XLM-R 使用 sentencepiece 模型进行文本标记化。下面，我们使用预训练的 sentencepiece 模型以及相应的词汇表来构建文本预处理管道，使用 torchtext 的 transforms。这些 transforms 通过 torchtext.transforms.Sequential()进行流水线处理，类似于 torch.nn.Sequential()，但是可以被 torchscript。请注意，transforms 支持批处理和非批处理文本输入，即可以传递单个句子或句子列表。

```py
import torchtext.transforms as T
from torch.hub import [load_state_dict_from_url](https://pytorch.org/docs/stable/hub.html#torch.hub.load_state_dict_from_url "torch.hub.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](https://pytorch.org/docs/stable/hub.html#torch.hub.load_state_dict_from_url "torch.hub.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](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader") 
```

或者我们也可以使用预训练模型附带的转换，它可以直接完成上述所有步骤

```py
text_transform = XLMR_BASE_ENCODER.transform() 
```

## 数据集

torchtext 提供了几个标准的 NLP 数据集。有关完整列表，请参阅文档：https://pytorch.org/text/stable/datasets.html。这些数据集是使用可组合的 torchdata datapipes 构建的，因此支持使用用户定义的函数和 transforms 进行标准的流控制和映射/转换。下面，我们演示如何使用文本和标签处理 transforms 来预处理 SST-2 数据集。

注意

使用 datapipes 目前仍然存在一些注意事项。如果您希望将此示例扩展到包括洗牌、多处理或分布式学习，请参阅此说明以获取进一步的指导。

```py
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](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.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](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader")(dev_datapipe, batch_size=None) 
```

或者我们也可以使用批处理 API（即对整个批次应用转换）

```py
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 提供了可以用于微调下游 NLP 任务的 SOTA 预训练模型。以下我们使用预训练的 XLM-R 编码器与标准基础架构，并附加一个分类器头来对其进行微调，以在 SST-2 二元分类任务上进行微调。我们将使用库中的标准分类器头，但用户可以定义自己适当的任务头并将其附加到预训练的编码器上。有关可用预训练模型的更多详细信息，请参阅文档：https://pytorch.org/text/main/models.html

```py
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) 
```

## 训练方法

现在让我们定义标准的优化器和训练标准，以及一些用于训练和评估的辅助函数。

```py
import torchtext.functional as F
from torch.optim import [AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW "torch.optim.AdamW")

learning_rate = 1e-5
optim = [AdamW](https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html#torch.optim.AdamW "torch.optim.AdamW")(model.parameters(), lr=learning_rate)
criteria = [nn.CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss "torch.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](https://pytorch.org/docs/stable/generated/torch.no_grad.html#torch.no_grad "torch.no_grad")():
        for batch in dev_dataloader:
            input = F.to_tensor(batch["token_ids"], padding_value=padding_idx).to(DEVICE)
            target = [torch.tensor](https://pytorch.org/docs/stable/generated/torch.tensor.html#torch.tensor "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 来处理数据分片。

```py
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](https://pytorch.org/docs/stable/generated/torch.tensor.html#torch.tensor "torch.tensor")(batch["target"]).to(DEVICE)
        train_step(input, target)

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

## 输出

```py
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 秒）

`下载 Python 源代码：sst2_classification_non_distributed.py`

`下载 Jupyter 笔记本：sst2_classification_non_distributed.ipynb`

[Sphinx-Gallery 生成的图库](https://sphinx-gallery.github.io)