Hugging Face NLP课程学习记录 - 3. 微调一个预训练模型

Hugging Face NLP课程学习记录 - 3. 微调一个预训练模型

说明：

• 首次发表日期：2024-12-11
• 官网： https://huggingface.co/learn/nlp-course/zh-CN/chapter3
• 关于： 阅读并记录一下，删除小部分内容，大多从原文摘录，润色一下原文

3. 微调一个预训练模型

安装Hugging Face的datasets模块：

pip install datasets

处理数据（Processing the data）

以下以一个批次（one batch）为例，说明了如何训练一个句子分类器（sequence classifier）：

import torch
from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification

# Same as before
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")

# This is new
batch["labels"] = torch.tensor([1,1])

optimizer = AdamW(model.parameters())
loss = model(**batch).loss
loss.backward()
optimizer.step()

当然，仅用两个句子训练模型不会产生非常好的结果。为了获得更好的结果，你需要准备一个更大的数据集。

在本节中，我们将使用MRPC（微软研究释义语料库）数据集（一个小数据集）作为示例。该数据集由5801对句子组成，每个句子对带有一个标签，指示它们是否为同义（即，如果两个句子的意思相同）。

从Hub加载数据（Loading a dataset from the Hub）

from datasets import load_dataset

raw_datasets = load_dataset("glue", "mrpc")
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、label 和 idx）以及不固定数量的行，这些行表示每个集合中的元素数量（因此，训练集中有 3,668 对句子，验证集中有 408 对句子，测试集中有 1,725 对句子）。

该命令会下载并缓存数据集，默认保存在 ~/.cache/huggingface/datasets 目录下。回想一下第2章的内容，你可以通过设置 HF_HOME 环境变量来自定义缓存文件夹。

我们可以像使用字典一样，通过索引访问 raw_datasets 对象中的每对句子：

raw_train_dataset = raw_datasets["train"]
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。这样可以查看每列的数据类型：

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

在幕后，label 是 ClassLabel 类型，整数到标签名称的映射存储在 names 字段中。0 对应于 not_equivalent，1 对应于 equivalent。

预处理数据集（Preprocessing a dataset）

为了预处理数据集，我们需要将文本转换为模型能够理解的数字。

from transformers import AutoTokenizer

checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
tokenized_sentence_1 = tokenizer(raw_datasets["train"]["sentence1"])
tokenized_sentence_2 = tokenizer(raw_datasets["train"]["sentence2"])

然而，我们不能仅仅传递两个序列给模型并期望获得关于这两句是否是释义的预测。我们需要将这两个序列作为一对来处理，并应用适当的预处理。幸运的是，分词器也可以接受一对序列，并按照 BERT 模型的预期方式进行准备：

inputs = tokenizer("This is the first sentence.", "This is the second one.")
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]
}

在这个例子中，类型标记ID(token_type_ids)的作用就是告诉模型输入的哪一部分是第一句，哪一部分是第二句。

如果我们将input_ids中的id转换回文字:

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_id 全部为 0，而对应 sentence2 [SEP] 的部分的 token_type_id 全部为 1。

请注意，如果你选择一个不同的检查点（checkpoint，意味着不同的模型），你的分词输入中不一定会有token_type_ids（例如，如果你使用DistilBERT模型，它们就不会被返回）。只有当模型知道如何处理它们时，才会返回它们，因为它在预训练期间已经见过它们。

---

在下一个句子预测任务中，会给模型输入成对的句子（带有随机遮罩的词元），并被要求预测第二个句子是否紧跟第一个句子。为了提高模型的泛化能力，数据集中一半的两个句子在原始文档中挨在一起，另一半的两个句子来自两个不同的文档。

一般来说，你不需要担心你的分词输入中是否有token_type_ids：只要你对分词器和模型使用相同的检查点，一切都会很好，因为分词器知道应该为它的模型提供什么。

预处理整个数据集：

tokenized_dataset = tokenizer(
    raw_datasets["train"]["sentence1"],
    raw_datasets["train"]["sentence2"],
    padding=True,
    truncation=True,
)

这方法效果很好，但有一个缺点是返回一个字典（包含我们的键：input_ids、attention_mask和token_type_ids，以及值是列表的列表）。此外，它只有在你有足够的RAM来存储整个数据集进行分词时才有效（而来自🤗 Datasets库的数据集是存储在磁盘上的Apache Arrow文件，因此你只需将请求的样本加载到内存中）。

为了将数据保存为数据集，我们将使用 Dataset.map() 方法。这也为我们提供了一些额外的灵活性，如果我们需要进行除了分词以外更多的预处理。map() 方法通过对数据集中的每个元素应用一个函数来工作，因此让我们定义一个对输入进行分词的函数：

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

这个函数接收一个字典（像我们数据集的项目）并返回一个新的字典，包含键input_ids、attention_mask和token_type_ids。注意，即使示例字典包含多个样本（每个键都是句子列表），它也能工作，因为分词器可以处理句子对的列表，就像之前看到的那样。这将允许我们在调用map()时使用batched=True选项，这将大大加快分词速度。

请注意，我们暂时没有在分词函数中包含padding参数。这是因为将所有样本填充到最大长度是低效的：更好的做法是在构建批次时对样本进行填充，这样我们只需要填充到该批次中的最大长度，而不是整个数据集中的最大长度。当输入长度变化很大时，这可以节省大量时间和处理能力！

以下是我们如何对整个数据集一次性应用分词函数的方式。在调用 map 方法时，我们使用了 batched=True，这样函数会一次作用于数据集的多个元素，而不是分别对每个元素单独处理。这种方式可以加快预处理速度。

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
tokenized_datasets

🤗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() 应用预处理函数时，你甚至可以通过传递 num_proc 参数来启用多进程处理。这里我们没有这么做，因为 🤗 Tokenizers 库已经使用了多线程来更快地对样本进行分词。

我们的 tokenize_function 返回一个包含 input_ids、attention_mask 和 token_type_ids 这三个键的字典，因此这三个字段会被添加到数据集的所有分区中。注意，如果预处理函数map()为现有键返回一个新值，那将会修改原有键的值。

最后一件我们需要做的事情是，当我们一起批处理元素时，将所有示例填充到最长元素的长度------我们称之为动态填充。

动态填充（Dynamic padding）

负责在批处理中将数据整理为一个batch的函数称为collate函数。它是你可以在构建DataLoader时传递的一个参数，默认是一个函数，它将把你的数据集转换为PyTorch张量，并将它们拼接起来(如果你的元素是列表、元组或字典，则会使用递归)。这在我们的这个例子中下是不可行的，因为我们的输入不是都是相同大小的。我们故意在之后每个batch上进行填充，避免有太多填充的过长的输入。这将大大加快训练速度，但请注意，如果你在TPU上训练，这可能会导致问题------TPU喜欢固定的形状，即使这需要额外的填充。

在实际操作中，我们需要定义一个 collate 函数，对想要批处理的数据集项应用适当的填充操作。幸运的是，🤗 Transformers 库通过 DataCollatorWithPadding 为我们提供了这样的函数。在实例化它时，你需要提供一个分词器（以便知道使用哪个填充标记，以及模型是否期望填充在输入的左侧或右侧），它将完成你需要做的一切：

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

为了测试这个新玩具，我们从训练集里选取一些样本，将它们打包成一个批次。在这里，我们移除了 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 微调模型（Fine-tuning a model with the Trainer API）

下面的示例假设您已经执行了上一节中的示例。下面这段代码，概括了您需要提前运行的代码：

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 在训练和评估中使用的所有超参数。您唯一必须提供的参数是保存训练模型的目录，以及训练过程中的检查点。对于其余的参数，您可以保留默认值，这对于基本微调应该非常有效。

from transformers import TrainingArguments

training_args = TrainingArguments("test-trainer")

第二步是定义我们的模型：

from transformers import AutoModelForSequenceClassification

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

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

一旦我们有了模型，就可以通过将之前构建的所有对象传递给它来定义一个 Trainer，包括模型本身、training_args、训练和验证数据集、data_collator 以及分词器：

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

请注意，当像我们这里一样传递分词器时，Trainer 默认使用的 data_collator 会是前面提到的 DataCollatorWithPadding，因此可以省略调用中的 data_collator=data_collator 这一行。

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

trainer.train()

这将开始微调（在 GPU 上大约需要几分钟），并每 500 步报告一次训练损失。不过，它不会告诉你模型的实际表现好坏。这是因为：

1. 我们没有通过设置 evaluation_strategy 为 "steps"（每隔 eval_steps 评估一次）或 "epoch"（每个 epoch 结束时评估一次）来指示 Trainer 在训练期间进行评估。
2. 我们没有为 Trainer 提供一个 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 x 2（408 是我们用于预测的数据集中的元素数量）。这些是我们传递给 predict() 的数据集每个元素的 logits（正如你在上一章看到的，所有 Transformer 模型返回的是 logits）。为了将它们转换为可以与标签进行比较的预测，我们需要取第二个轴上最大值的索引：

import numpy as np

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

现在我们可以将这些预测值与标签进行比较。为了构建我们的 compute_metrics() 函数，我们将依赖于 🤗 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}

将所有内容整合在一起，我们得到了我们的 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()

Trainer 可以开箱即用地支持多 GPU 或 TPU，并提供许多选项，例如混合精度训练（在训练参数中设置 fp16 = True）。

一个完整的训练 （A full training）

现在，我们来看如何在不使用 Trainer 类的情况下达到与上一节相同的效果。同样，我们假设你已经完成了第 2 节中的数据处理。以下是涵盖所有所需内容的简要总结：

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)

训练前的准备 （Preparing for training）

在编写训练循环之前，我们需要定义一些对象。首先是数据加载器（dataloaders），它们将用于批量迭代数据。但在定义这些加载器之前，我们需要对 tokenized_datasets 进行一些后处理，以处理一些 Trainer 自动为我们完成的任务。具体来说，我们需要：

1. 移除模型不需要的列（例如 sentence1 和 sentence2 列）。
2. 将 label 列重命名为 labels（因为模型期望参数名为 labels）。
3. 设置数据集的格式，使其返回 PyTorch 张量而不是列表。

我们的 tokenized_datasets 为上述每一步提供了一个方法：

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"]

至此，我们可以轻松定义数据加载器:

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
{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])}

请注意，由于我们在训练数据加载器中设置了 shuffle=True，并在批次内填充到最大长度，实际的形状可能会与你的结果略有不同。

现在我们已经完全完成了数据预处理（对任何机器学习从业者来说，这都是一个令人满足却难以实现的目标），接下来我们来看模型。我们像上一节一样实例化模型：

from transformers import AutoModelForSequenceClassification

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

为了确保训练过程中一切顺利，我们将批次数据传递给这个模型：

outputs = model(**batch)
print(outputs.loss, outputs.logits.shape)

tensor(0.5441, grad_fn=<NllLossBackward>) torch.Size([8, 2])

所有 🤗 Transformers 模型在提供标签时都会返回损失值，我们还会得到 logits（每个输入对应两个值，因此是一个大小为 8 x 2 的张量）。

我们几乎已经准备好编写训练循环了！只差两样东西：优化器和学习率调度器。由于我们尝试手动复现 Trainer 的默认行为，因此会使用相同的默认设置。Trainer 使用的优化器是 AdamW，它与 Adam 类似，但在权重衰减正则化上有所改进（参见 Ilya Loshchilov 和 Frank Hutter 的论文 Decoupled Weight Decay Regularization）：

from transformers import AdamW

optimizer = AdamW(model.parameters(), lr=5e-5)

最后，默认使用的学习率调度器是从最大值（5e-5）线性衰减到 0。为了正确地定义它，我们需要知道训练步骤的数量，这个数量是我们想要运行的 epoch 数与训练批次数量的乘积（训练数据加载器的长度）。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)

训练循环（The training loop）

最后一件事：如果有 GPU 可用，我们会希望使用它（在 CPU 上训练可能需要几个小时，而不是几分钟）。为此，我们需要定义一个设备，将模型和批数据放到这个设备上：

import torch

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
device

device(type='cuda')

我们现在准备好训练了！为了了解训练何时结束，我们使用 tqdm 库,在训练步骤数上添加了一个进度条:

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)

你可以看到，训练循环的核心与引言中的非常相似。我们没有要求任何报告，因此这个训练循环不会告诉我们模型的表现如何。我们需要添加一个评估循环来获取这些信息。

评估循环 （The evaluation loop）

正如我们之前所做的，我们将使用 🤗 Evaluate 库提供的一个指标（metric）。我们已经见过了 metric.compute() 方法，但实际上，当我们在预测循环中使用 add_batch() 方法时，指标可以为我们累积批次数据。一旦我们累积了所有批次，我们可以使用 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()

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

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 上使用 TPU 进行测试），只需将代码粘贴到 training_function() 并使用以下命令运行最后一个单元格:

from accelerate import notebook_launcher

notebook_launcher(training_function)

你可以在🤗 Accelerate repo找到更多的示例。