NLP 课程文档

使用Trainer API微调模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验

开始使用

使用Trainer API微调模型

Ask a Question Open In Colab Open In Studio Lab

🤗 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)

训练

在定义 Trainer 之前,第一步是定义一个 TrainingArguments 类,其中将包含 Trainer 用于训练和评估的所有超参数。您必须提供的唯一参数是训练好的模型将保存到的目录,以及保存过程中的检查点。对于所有其他参数,您可以保留默认值,这对于基本微调应该可以很好地工作。

from transformers import TrainingArguments

training_args = TrainingArguments("test-trainer")

💡 如果您希望在训练期间自动将模型上传到 Hub,请在 TrainingArguments 中传递 push_to_hub=True。我们将在 第 4 章 中了解更多关于此内容。

第二步是定义我们的模型。与 上一章 一样,我们将使用 AutoModelForSequenceClassification 类,有两个标签。

from transformers import AutoModelForSequenceClassification

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

您会注意到,与 第 2 章 不同,在实例化此预训练模型后,您会收到一个警告。这是因为 BERT 没有在对句子对进行分类上进行预训练,因此预训练模型的头已被丢弃,并且已添加了一个适合序列分类的新头。警告表明一些权重未被使用(对应于已丢弃的预训练头的权重),而另一些权重则被随机初始化(新头的权重)。最后,它鼓励您训练模型,这正是我们现在要做的。

拥有模型后,我们可以通过将所有构造好的对象传递给它来定义 Trainer——modeltraining_args、训练和验证数据集、我们的 data_collator 和我们的 tokenizer

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 节中向您展示处理的这一部分仍然很重要!

要对我们的数据集微调模型,我们只需要调用 Trainertrain() 方法。

trainer.train()

这将开始微调(在 GPU 上应该需要几分钟)并在每 500 步报告训练损失。但是,它不会告诉您模型的性能如何(或有多差)。这是因为

  1. 我们没有告诉 Trainer 在训练期间进行评估,方法是将 evaluation_strategy 设置为 "steps"(每 eval_steps 步评估一次)或 "epoch"(在每个 epoch 结束时评估)。
  2. 我们没有为 Trainer 提供 compute_metrics() 函数来在上述评估期间计算指标(否则评估只会打印损失,这并不是一个非常直观的数字)。

评估

让我们看看如何构建一个有用的 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() 方法的输出是另一个名为元组,具有三个字段:predictionslabel_idsmetricsmetrics 字段将只包含传递给数据集的损失,以及一些时间指标(预测花费的时间,总时间和平均时间)。一旦我们完成了 compute_metrics() 函数并将其传递给 Trainer,该字段还将包含 compute_metrics() 返回的指标。

如您所见,predictions 是一个二维数组,形状为 408 x 2(408 是我们使用的数据集中元素的数量)。这些是我们传递给 predict() 的数据集每个元素的 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。这是用于评估 GLUE 基准测试中 MRPC 数据集结果的两个指标。 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)

为了看到它在每个 epoch 结束时报告指标的实际应用,以下是如何使用此 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)。我们将在第 10 章中介绍它支持的所有内容。

这结束了使用 Trainer API 进行微调的介绍。在第 7 章中将提供针对大多数常见 NLP 任务的示例,但现在让我们看看如何在纯 PyTorch 中执行相同的操作。

✏️ 动手试试! 使用你在第 2 节中完成的数据处理,在 GLUE SST-2 数据集上微调模型。