huggingface NLP工具包教程3：微调预训练模型
huggingface NLP工具包教程3：微调预训练模型
引言
在上一章我们已经介绍了如何使用 tokenizer 以及如何使用预训练的模型来进行预测。本章将介绍如何在自己的数据集上微调一个预训练的模型。在本章，你将学到：
- 如何从 Hub 准备大型数据集
- 如何使用高层 Trainer API 微调模型
- 如何使用自定义训练循环
- 如何利用 Accelerate 库，进行分布式训练
如果想要将将经过训练的权重上传到 Hugging Face Hub，需要注册一个 huggingface.co 账号：创建账号。
处理数据
以下是如何训练一个序列分类器（以一个 batch 为例）：
import torch
from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification
# 与之前章节一致
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
sequences = [
"I've been waiting for a HuggingFace course my whole life.",
"This course is amazing!",
]
batch = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")
# 以下是训练部分
batch["labels"] = torch.tensor([1, 1])
optimizer = AdamW(model.parameters())
loss = model(**batch).loss
loss.backward()
optimizer.step()
当然，向上面这样仅使用两个句子进行训练肯定无法得到满意的结果。我们需要一个更大的数据集。
本节我们将以 MRPC (Microsoft Research Paraphrase Corpus) 数据集为例，该数据集由 5801 对句子组成，并带有一个标签，表明它们是否是转述（即，如果两个句子的意思相同）。本章选择该数据集，因为它是一个小数据集，所以很容易进行训练。
从Hub中加载数据
Hub 不仅包括模型，它同样包含不同语言的多个数据集。可以在这里查看数据集，推荐读者在过完本节之后自己试着加载并处理一个新数据集，文档参考这里。不过现在，让我们先来看 MRPC 数据集，他是组成 GLUE benchmark 的十个数据集之一。GLUE benchmark 是一个在 10 个不同的文本分类任务上评估机器学习模型的学术基准。
Datasets 库提供了一些非常简单的命令来下载并缓存 Hub 中的数据集。例如下载 MRPC 数据集：
from datasets import load_dataset
raw_datasets = load_dataset("glue", "mrpc")
print(raw_datasets)
# 输出：
DatasetDict({
train: Dataset({
features: ['sentence1', 'sentence2', 'label', 'idx'],
num_rows: 3668
})
validation: Dataset({
features: ['sentence1', 'sentence2', 'label', 'idx'],
num_rows: 408
})
test: Dataset({
features: ['sentence1', 'sentence2', 'label', 'idx'],
num_rows: 1725
})
})
可以看到，我们得到了一个 DatasetDict 对象，其中包含了训练集、验证集和测试集。其中每个又包括了几个列（sentence1,sentence2,labe,idx），和一个行数值，表示每个集合中的样本个数。
上述命令会下载并缓存指定数据集，默认缓存目录是 ~/.cache/huggingface/datasets。同样可以通过环境变量 HF_HOME 来修改。
我们可以通过索引来访问 raw_datasets 中的每对句子：
raw_train_dataset = raw_datasets["train"]
print(raw_train_dataset[0])
# 输出：
{'idx': 0,
'label': 1,
'sentence1': 'Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .',
'sentence2': 'Referring to him as only " the witness " , Amrozi accused his brother of deliberately distorting his evidence .'}
可以看到标签已经是整型了，所以这里不再需要进一步处理。如果想要知道哪个整型值对应哪个标签，可以查看 raw_train_dataset 的 features 属性。它会返回每一列的类型：
print(raw_train_dataset.features)
# 输出：
{'sentence1': Value(dtype='string', id=None),
'sentence2': Value(dtype='string', id=None),
'label': ClassLabel(num_classes=2, names=['not_equivalent', 'equivalent'], names_file=None, id=None),
'idx': Value(dtype='int32', id=None)}
可以看到，标签的类型为 ClassLabel，整数到标签名称的映射存储在 names folder 中。0 对应于 not_equivalent，1 对应于 equivalent。
数据集预处理
我们需要将原文本数据转换为模型可以接收的数值型，之前的章节已经介绍过，这由 tokenizer 完成。tokenizer 既可以接收一个句子，也可以接收一组句子。因此，我们可以直接将数据集中每个句子对中的 ”第一个句子“ 和 ”第二个句子” 直接送入给 tokenizer：
from transformers import AutoTokenizer
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
tokenized_sentences_1 = tokenizer(raw_datasets["train"]["sentence1"])
tokenized_sentences_2 = tokenizer(raw_datasets["train"]["sentence2"])
然而，我们不能仅仅是直接将两个序列传给模型，然后让模型预测两个序列是否是释义关系。我们需要将两个序列处理成一个序列对，并进行适当的预处理。幸运的是，强大的 tokenizer 也可以接收一对序列并将它们处理成 BERT 模型需要的形式：
inputs = tokenizer("This is the first sentence.", "This is the second one.")
print(inputs)
# 输出：
{
'input_ids': [101, 2023, 2003, 1996, 2034, 6251, 1012, 102, 2023, 2003, 1996, 2117, 2028, 1012, 102],
'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
}
输出中的 inpud_ids 和 attention_mask 我们在前一章已经介绍过，但是 token_type_ids 当时没有提，在这里，它负责告诉模型哪部分输入是第一个句子，哪部分是第二个句子。
我们可以对 input_ids 进行解码会单词：
tokenizer.convert_ids_to_tokens(inputs["input_ids"])
# 输出：
['[CLS]', 'this', 'is', 'the', 'first', 'sentence', '.', '[SEP]', 'this', 'is', 'the', 'second', 'one', '.', '[SEP]']
可以看到，模型期望的输入是将两个句子合并成 [CLS] sentence1 [SEP] sentence2 [SEP] 的形式，这与我们给出的 token_type_ids 是对齐的：
['[CLS]', 'this', 'is', 'the', 'first', 'sentence', '.', '[SEP]', 'this', 'is', 'the', 'second', 'one', '.', '[SEP]']
[ 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
第一个句子 [CLS] sentence1 [SEP] 对应的 token_type_ids 都为 0，而第二个句子 sentence2 [SEP] 对应的 token_type_ids 为 1。
注意 token_id_types 并不是所有模型都必须的，只有模型预训练任务中需要这种输入时才需要。也就是说如果我们用的是其他预训练模型（比如 DistilBERT），可能就不需要 token_id_types，这时 tokenizer 也不会返回该键。
这里的 BERT 模型预训练时是需要 token_type_ids 的，BERT 模型在预训练时，除了第一章中介绍过的 masked language modeling 任务之外，还有另一个预训练任务：next sentence prediction ，该任务的目标是建模一对句子之间的关系。具体来说，该任务需要预测一对句子（当然也包括一些 mask 掉的 token）之中，第二句在语义上是否是第一句的后一句。为了使得该任务有一定难度，一半训练数据中两个句子来自同一个文本（是或不是连续两句），另一半训练数据则来自不同文本（肯定不是连续两句）。
总之，只要保证加载的模型和加载的 tokenizer 是来自同一个预训练权重，我们就不太需要担心输入数据中是否需要包含 token_type_ids ，因为 tokenizer 会帮我们把一切都准备好。
我们已经了解如何使用 tokenizer 来处理一对句子，现在用它来处理整个数据集：直接将一组句子对送入到 tokenizer 中，并且指定填充和截断：
tokenized_dataset = tokenizer(
raw_datasets["train"]["sentence1"],
raw_datasets["train"]["sentence2"],
padding=True,
truncation=True,
)
这样可以正常工作，这样的缺点是它会直接返回一个字典（包括 input_ids, attention_mask, and token_type_ids），这需要我们的机器有足够大的内存来存储整个数据集。
我们会使用 Dataset.map() 方法来封装预处理过程，该方法比较灵活，方便除了 tokenization 之外，在预处理阶段添加更多操作。map() 方法的工作方式是将一个函数内定义的操作施加到 Dataset 中的每个元素上。这里我们先定义一个处理函数：
def tokenize_function(example):
return tokenizer(example["sentence1"], example["sentence2"], truncation=True)
该函数接收一个字典（比如数据集中的一个样本）作为输入，并返回一个新的字典，包含 input_ids, attention_mask, and token_type_ids 这几个键。如果 example 字典包含多个样本（每个键都是一个句子列表），也是可行的。因为前面介绍过，tokenizer 可以处理成对的句子列表。我们在调用 map() 的时传入参数 batched=True，这将大大加快 tokenization 过程。tokenizer 来自由 Rust 编写的 Tokenizer 库。
注意我们并没有在 tokenize_function 中设置 padding 参数，因为直接按照全数据集的最大长度进行填充是很低效的，最好是在构建 batch 时进行填充，这样就只需要按照 batch 内的最大长度进行填充即可，而不需要按照整个数据集的最大长度进行填充。
我们通过设置 batched=True 来同时处理多个元素，而非一个一个处理，这会使得整个预处理过程更快：
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
print(tokenized_datasets)
# 输出：
DatasetDict({
train: Dataset({
features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],
num_rows: 3668
})
validation: Dataset({
features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],
num_rows: 408
})
test: Dataset({
features: ['attention_mask', 'idx', 'input_ids', 'label', 'sentence1', 'sentence2', 'token_type_ids'],
num_rows: 1725
})
})
可以看到，map 会通过在字典中添加处理后得到的新键来完成预处理过程。如果某些处理结果会得到已有键的新值，map 也会用新值覆盖掉旧值。
我们还可以在调用 map 方法时，通过传入 num_proc 参数来进行多进程的数据预处理。上例中没有这么做是因为 tokenizer 库已经使用了多线程来进行加速，但如果在实际中没有使用该库的 tokenizer 的话，可以在这里使用多进程进行加速。
最后一件要介绍的事情是动态填充（dynamic padding），它负责将一个 batch 内的所有序列填充到该 batch 内原始序列的最长长度。
动态填充
collate function 负责将一个 batch 内所有的样本放到一起。在 Pytorch 中，它是我们构建 DataLoader 时一个可选的参数，默认的 collate function 会简单地将所有的样本数据转换为张量并拼接在一起。这肯定是不能直接用的，因为目前我们还没有进行填充，batch 内的每个样本不是等长的。之前我们故意先没有填充，因为我们想在构建每个 batch 时进行填充，避免一个 batch 内无意义的填充。这将大大加快训练速度，但请注意，如果时在TPU上训练，可能会导致问题，因为 TPU 更偏好固定的形状，即使这需要额外的填充。
实际中，我们会定义一个 collate function，来对每个 batch 内的样本进行填充。Transformers 库通过 DataCollatorWithPadding 提供这样的功能。在对其进行实例化时，传入我们的 tokenizer，因为它需要知道 padding token 是什么，并且要知道该模型需要再输入序列的左侧填充还是右侧填充，之后，它会为我们做好一切。
from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
我们从我们的训练集中抓取一些样本来将它们打包成 batch。这里，我们删除列idx、sentence1 和 sentence2，因为不需要它们而且它们包含字符串（我们不能用字符串创建张量），并查看批次中每个元素的长度：
samples = tokenized_datasets["train"][:8]
samples = {k: v for k, v in samples.items() if k not in ["idx", "sentence1", "sentence2"]}
[len(x) for x in samples["input_ids"]]
# 输出：
[50, 59, 47, 67, 59, 50, 62, 32]
毫无疑问，一个批次中的原始序列长度是不一致的，从 32 到 67。动态填充意味着该批次内的序列都要填充到序列中的最长长度为 67。如果不使用动态填充的话，这些样本的序列长度都要被填充到整个数据集的最大长度，或者模型能够接收的最大长度。下面我们再检查一下 data_collator 是否正确进行了动态填充：
batch = data_collator(samples)
{k: v.shape for k, v in batch.items()}
# 输出：
{'attention_mask': torch.Size([8, 67]),
'input_ids': torch.Size([8, 67]),
'token_type_ids': torch.Size([8, 67]),
'labels': torch.Size([8])}
结果一切正常。至此，我们就完成了对原始文本数据的预处理，准备正式开始进行微调。
使用Trainer API对模型进行微调
Transformers 库提供了一个 Trainer 类来帮助用户在自己的数据集上对预训练模型进行微调。在完成上一节的数据预处理准备之后，马上就可以开始微调训练了。
下面的代码是上一节数据预处理的汇总：
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")
第二步就是要定义模型。这里我们使用 AutoModelForSequenceClassification，类别数为 2：
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
这一步实例化一个模型之后，会得到一个警告。这是因为 BERT 模型预训练任务不是对句子对进行分类，所以预训练时的模型头部被直接丢掉，然后换用一个新的适合指定任务的头部。该警告表明预训练模型中有一部分权重（此处就是模型头部部分）没有用到，并且另外有一些权重（此处即新的头部）是随机初始化的。
当我们定义好模型之后，就可以定义 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 这一行是可以跳过的。
接下来，直接调用 trainer.train() 方法就可以开始微调模型：
trainer.train()
这就会开始微调，并每过 500 个 steps 就报告一次损失。但是这并不能告诉我们模型实际性能如何，因为：
- 我们没有通过设置
evaluation_strategy来告诉模型在每个 step 或每个 epoch 之后对模型进行评估 - 我们没有提供
compute_metrics()函数给模型，来告诉他如何计算指标
evaluation
接下来介绍如何构建一个有用的 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() 返回的指标。
可以看到 predictions 是一个二维数组，形状为 408 × 2 408\times 2 408×2 （ 408 是数据集中的样本个数）。这是数据集中每个样本预测结果 logits，我们需要取它们最大值的索引，来得到模型最终预测的类别：
import numpy as np
preds = np.argmax(predictions.predictions, axis=-1)
得到最终的预测类别之后，就可以与标签进行对比，计算指标。我们基于 Evaluate 库来构建 compute_metric() 函数。加载 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 Benchmark 上评测 MRPC 数据集所用的指标。在 BERT 原论文中报告的结果中，base 模型的 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 结束时查看这些指标，我们重新定义一个 Trainer，将 compute_metrics 函数加进来：
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='epoch'，并创建了新的模型。启动训练：
trainer.train()
这一次，它将在每个 epoch 结束时，除了会报告训练损失，还会报告验证损失和和指标。同样，由于模型的随机头部初始化，这次达到的准确率和F1分数可能与之前有所不同，但差别不会太大。
Trainer 可在多个 GPU 或 TPU 上开箱即用，并提供许多选项，如混合精度训练（在训练参数中使用 fp16=True）。我们将在第10章介绍更多。
使用Trainer API进行微调的介绍到此结束。第7章将给出为最常见的NLP任务执行此操作的示例，现在让我们先看看如何在纯 PyTorch 中执行相同的操作。
完整训练
现在我们将看到如何在不使用 Trainer 类，而使用纯 Pytorch，做到与上一节相同的事情。再次回顾第二节数据处理如下：
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)
准备训练
在实际编写训练脚本之前，我们需要先定义几个对象。第一个是我们将用于迭代batch 数据加载程序。但在定义这些 DataLoader 之前，我们需要对经过 tokenize 的数据集进行一些后处理，在上一节中 Trainer 自动为我们做了这些事情。具体来说，我们需要：
- 删除与模型不需要的值相对应的列（如句子1和句子2列）。
- 将列 label 重命名为 labels（因为模型需要的对应参数名为 labels）。
- 设置数据集的格式，使得它们返回 PyTorch 张量而不是列表。
tokenized_dataset 为每个步骤提供了对应方法：
tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")
tokenized_datasets["train"].column_names
然后检查一下结果是否对应我们模型需要的键：
["attention_mask", "input_ids", "labels", "token_type_ids"]
现在准备工作做好了，开始定义 dataloader：
from torch.utils.data import DataLoader
train_dataloader = DataLoader(
tokenized_datasets["train"], shuffle=True, batch_size=8, collate_fn=data_collator
)
eval_dataloader = DataLoader(
tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator
)
可以用过以下方法快速检查数据加载过程没有错误：
for batch in train_dataloader:
break
print({k: v.shape for k, v in batch.items()})
# 输出：
{'attention_mask': torch.Size([8, 65]),
'input_ids': torch.Size([8, 65]),
'labels': torch.Size([8]),
'token_type_ids': torch.Size([8, 65])}
训练数据的 Dataloader 设置了 shuffle=True，并且在 batch 中填充了最大长度，因此每个人实际查看的形状可能会略有不同。
现在已经完全完成了数据预处理，开始准备模型。我们与上一节中所做的完全一样进行实例化：
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
传入一个 batch 测试有无问题：
outputs = model(**batch)
print(outputs.loss, outputs.logits.shape)
# 输出：
tensor(0.5441, grad_fn=<NllLossBackward>) torch.Size([8, 2])
当提供标签时，所有 Transformer 模型都会返回损失，也会得到 logits（在我们的 batch 中，每次输入两个，所以张量大小为 8 x 2 8 x 2 8x2）。
还差两个东西：优化器（optimizer）和学习率调度器（learning rate scheduler）。由于我们试图以手动操作复现 trainer 的结果，因此这里使用相同的默认值。Trainer 使用的优化器是AdamW，它与Adam 相同，但对权重衰减正则项进行了扭曲（参见“Decoupled Weight Decay Regularization” ）：
from transformers import AdamW
optimizer = AdamW(model.parameters(), lr=5e-5)
最后，默认情况下使用的学习速率调度器是从最大值（5e-5）到 0 的线性衰减。我们需要知道总训练步数，即我们要运行的 epoch 数乘以训练 batch 数（即 DataLoader 的长度）。trainer 默认使用三个 epoch，因此我们按照如下定义：
from transformers import get_scheduler
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=num_training_steps,
)
print(num_training_steps)
# 输出：
1377
训练循环
然后我们再定义一下训练所用的设备：
import torch
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
print(device)
# 输出：
device(type='cuda')
现在可以编写训练脚本并进行训练了：
from tqdm.auto import tqdm
progress_bar = tqdm(range(num_training_steps))
model.train()
for epoch in range(num_epochs):
for batch in train_dataloader:
batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
可以看到，训练循环的核心步骤看起来很像引言中。我们没有打印任何指标，所以这个训练循环不会告诉我们关于模型如何运行的任何信息。我们需要为此添加一个评估循环。
评估循环
与前面一样，我们将使用 Evaluate 库。我们介绍过了 metric.compute() 方法：
import evaluate
metric = evaluate.load("glue", "mrpc")
model.eval()
for batch in eval_dataloader:
batch = {k: v.to(device) for k, v in batch.items()}
with torch.no_grad():
outputs = model(**batch)
logits = outputs.logits
predictions = torch.argmax(logits, dim=-1)
metric.add_batch(predictions=predictions, references=batch["labels"])
metric.compute()
# 输出：
{'accuracy': 0.8431372549019608, 'f1': 0.8907849829351535}
同样地，结果会有随机性，但应该差不太多。
使用accelerate库加速训练
之前的训练脚本是工作在单个 CPU 或单个 GPU 上的，通过使用 accelerate 库 ，只需少量改动就可以运行在多 GPU/TPU 上。从创建训练/验证数据集开始，以下是手动训练循环的代码：
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=num_training_steps,
)
progress_bar = tqdm(range(num_training_steps))
model.train()
for epoch in range(num_epochs):
for batch in train_dataloader:
batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
以下是改为 accelerate 加速所改动的部分：
+ from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler
+ accelerator = Accelerator()
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)
- device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
- model.to(device)
+ train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
+ train_dataloader, eval_dataloader, model, optimizer
+ )
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=num_training_steps
)
progress_bar = tqdm(range(num_training_steps))
model.train()
for epoch in range(num_epochs):
for batch in train_dataloader:
- batch = {k: v.to(device) for k, v in batch.items()}
outputs = model(**batch)
loss = outputs.loss
- loss.backward()
+ accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
添加的第一行是导包。第二行实例化一个 Accelerator 对象，该对象将查看环境并初始化分布式设置。Accelerate 库会处理数据在设备上的存放，因此可以删除将模型放置在设备上的那一行（或者，也可以将 device 改为 accelerator.device）。
然后，将数据加载器、模型和优化器送入到 accelerator.prepare()。这将把这些对象包装在适当的容器中，以确保分布式训练正常工作。最后一处的改动是删除将 batch 放在设备上的那一行（同样，如果想保留该行，可以将其更改为使用 accelerator.device 并将 loss.backward() 改为 accelerator.backward(loss) 。
以下是完成的使用 accelerate 库加速的代码：
from accelerate import Accelerator
from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler
accelerator = Accelerator()
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=3e-5)
train_dl, eval_dl, model, optimizer = accelerator.prepare(
train_dataloader, eval_dataloader, model, optimizer
)
num_epochs = 3
num_training_steps = num_epochs * len(train_dl)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=num_training_steps,
)
progress_bar = tqdm(range(num_training_steps))
model.train()
for epoch in range(num_epochs):
for batch in train_dl:
outputs = model(**batch)
loss = outputs.loss
accelerator.backward(loss)
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
将上述代码放在 train.py 脚本中，该脚本可以在任何类型的分布式设置上运行。如果要测试自己的分布式设置，运行：
accelerate config
根据提示进行一些配置，并保存到配置文件中，然后运行：
accelerate launch train.py
这就将启动分布式训练。
如果想要在 Notebook （比如 Colab ）中运行，只需将上述代码封装到 training_function() 中，然后运行：
from accelerate import notebook_launcher
notebook_launcher(training_function)
更多示例可参考：Accelerate repo.
总结
概括一下，在本章中，我们：
-
了解 Hub 中的 datasets
-
学习了如何加载和预处理数据集，包括使用动态填充和 collator
-
实现了自己对模型的微调和评估
-
实现了更底层的（基于纯 Pytorch）的训练循环
-
使用 accelerate 进行分布式训练