使用 Trainer API 微调模型

🤗 Transformers 提供了一个 Trainer 类，可以帮助你在数据集上微调任何预训练模型。在上一节中完成所有数据预处理工作后，你只需完成几个步骤来定义 Trainer 。最困难的部分可能是准备运行 Trainer.train() 所需的环境，因为在 CPU 上运行速度非常慢。如果你没有设置 GPU，可以使用 Google Colab （国内网络无法使用） 上获得免费的 GPU 或 TPU。

下面的示例假设你已经执行了上一节中的示例。下面是在开始学习这一节之前你需要运行的代码：

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding

raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)


def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)


tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

Training

在我们定义 Trainer 之前第一步要定义一个 TrainingArguments 类，它包含 Trainer 在训练和评估中使用的所有超参数。你只需要提供的参数是一个用于保存训练后的模型以及训练过程中的 checkpoint 的目录。对于其余的参数你可以保留默认值，这对于简单的微调应该效果就很好了。

from transformers import TrainingArguments

training_args = TrainingArguments("test-trainer")

第二步是定义我们的模型。与 前一章 一样，我们将使用 AutoModelForSequenceClassification 类，它有两个参数：

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

你会注意到，和 第二章 不同的是，在实例化此预训练模型后会收到警告。这是因为 BERT 没有在句子对分类方面进行过预训练，所以预训练模型的 head 已经被丢弃，而是添加了一个适合句子序列分类的新头部。这些警告表明一些权重没有使用（对应于被放弃的预训练头的权重），而有些权重被随机初始化（对应于新 head 的权重）。

一旦有了我们的模型，我们就可以定义一个 Trainer 把到目前为止构建的所有对象 —— model ，training_args，训练和验证数据集， data_collator 和 tokenizer 传递给 Trainer ：

from transformers import Trainer

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

请注意，当你在这里传递 tokenizer 时， Trainer 默认使用的 data_collator 是之前预定义的 DataCollatorWithPadding。所以你可以在本例中可以跳过 data_collator=data_collator 一行。在第 2 节中向你展示这部分处理过程仍然很重要！

要在我们的数据集上微调模型，我们只需调用 Trainer 的 train() 方法：

trainer.train()

开始微调（在 GPU 上应该需要几分钟），每 500 步报告一次训练损失。然而它不会告诉你模型的性能（或质量）如何。这是因为：

1. 我们没有告诉 Trainer 在训练过程中进行评估，比如将 evaluation_strategy 设置为“ step ”（在每个 eval_steps 步骤评估一次）或“ epoch ”（在每个 epoch 结束时评估）。
2. 我们没有为 Trainer 提供一个 compute_metrics() 函数来计算上述评估过程的指标（否则评估将只会输出 loss，但这不是一个非常直观的数字）。

评估

让我们看看如何构建一个有用的 compute_metrics() 函数，并在下次训练时使用它。该函数必须接收一个 EvalPrediction 对象（它是一个带有 predictions 和 label_ids 字段的参数元组），并将返回一个字符串映射到浮点数的字典（字符串是返回的指标名称，而浮点数是其值）。为了从我们的模型中获得预测结果，可以使用 Trainer.predict() 命令：

predictions = trainer.predict(tokenized_datasets["validation"])
print(predictions.predictions.shape, predictions.label_ids.shape)

(408, 2) (408,)

predict() 方法的输出另一个带有三个字段的命名元组: predictions label_ids 和 metrics metrics 字段将只包含所传递的数据集的损失,以及一些时间指标(总共花费的时间和平均预测时间)。当我们定义了自己的 compute_metrics() 函数并将其传递给 Trainer 该字段还将包含 compute_metrics() 返回的结果。predict() 是一个二维数组,形状为 408 × 2(408 是我们使用的数据集中的元素数量),这是我们传递给 pprdict() 的数据集中每个元素的 logits(正如在 前一章 中看到的,所有 Transformer 模型都返回 logits)。为了将它们转化为可以与我们的标签进行比较的预测值,我们需要找出第二个维度上取值最大的索引:

import numpy as np

preds = np.argmax(predictions.predictions, axis=-1)

我们现在可以将这些 preds 与标签进行比较。为了构建我们的 compute_metric() 函数，我们将使用 🤗 Evaluate 库中的指标。我们可以像加载数据集一样轻松地加载与 MRPC 数据集关联的指标，这次是使用 evaluate.load() 函数。返回的对象有一个 compute() 方法，我们可以用它来进行指标的计算：

import evaluate

metric = evaluate.load("glue", "mrpc")
metric.compute(predictions=preds, references=predictions.label_ids)

{'accuracy': 0.8578431372549019, 'f1': 0.8996539792387542}

你得到的确切结果可能会有所不同，因为模型头部的随机初始化可能会改变其指标。在这里，我们可以看到我们的模型在验证集上的准确率为 85.78％，F1 分数为 89.97。这是用于评估 MRPC 数据集在 GLUE 基准测试中的结果的两个指标。在 BERT 论文中的表格中，基础模型的 F1 分数为 88.9。那是 uncased 模型，而我们现在正在使用 cased 模型，这解释了为什么我们得到了更好的结果。

最后把所有东西打包在一起，我们就得到了 compute_metrics() 函数：

def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

为了查看模型在每个训练周期结束时的好坏，下面是我们如何使用 compute_metrics() 函数定义一个新的 Trainer

training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

请注意，我们设置了一个新的 TrainingArguments ，其 evaluation_strategy 设置为 epoch 并且创建了一个新模型。如果不创建新的模型就直接训练，就只会继续训练我们已经训练过的模型。为了启动新的训练，我们执行：

trainer.train()

这一次，它将在每个 epoch 结束时在训练损失的基础上报告验证损失和指标。同样，由于模型的随机头部初始化，达到的准确率/F1 分数可能与我们发现的略有不同，这是由于模型头部的随机初始化造成的，但应该相差不多。 Trainer 可以在多个 GPU 或 TPU 上运行，并提供许多选项，例如混合精度训练（在训练的参数中使用 fp16 = True ）。我们将在第十章讨论它支持的所有内容。

使用 Trainer API 微调的介绍到此结束。在 第七章 中会给出一个对大多数常见的 NLP 任务进行训练的例子，但现在让我们看看如何在 PyTorch 中做相同的操作。