NLP 课程文档

完整训练

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

完整训练

Ask a Question Open In Colab Open In Studio Lab

现在,我们将了解如何实现与上一节相同的结果,而无需使用 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)

准备训练

在实际编写训练循环之前,我们需要定义一些对象。第一个是我们将用于迭代批次的数据加载器。但在我们能够定义这些数据加载器之前,我们需要对我们的 tokenized_datasets 应用一些后处理,以解决 Trainer 自动为我们完成的一些事项。具体来说,我们需要

  • 删除与模型不期望的值相对应的列(例如 sentence1sentence2 列)。
  • label 列重命名为 labels(因为模型期望参数名为 labels)。
  • 设置数据集的格式,以便它们返回 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 模型在提供 labels 时都会返回损失,我们还会获得 logits(对于批次中的每个输入,都有两个,因此张量的尺寸为 8 x 2)。

我们几乎准备好编写训练循环了!我们只需要两样东西:优化器和学习率调度器。由于我们正在尝试手动复制 Trainer 所做的事情,因此我们将使用相同的默认值。Trainer 使用的优化器是 AdamW,它与 Adam 相同,但对于权重衰减正则化而言略有不同(参见 “解耦权重衰减正则化”,由 Ilya Loshchilov 和 Frank Hutter 撰写)

from transformers import AdamW

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

最后,默认情况下使用的学习率调度器只是一个从最大值(5e-5)到 0 的线性衰减。为了正确定义它,我们需要知道将要执行的训练步骤数量,即我们要运行的时期数量乘以训练批次数量(即训练数据加载器的长度)。Trainer 默认使用三个时期,因此我们将遵循此规则

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

训练循环

最后一点:如果我们有权访问 GPU,则需要使用它(在 CPU 上,训练可能需要几个小时而不是几分钟)。为此,我们定义一个 device,我们将模型和批次放到其中

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)

您可以看到,训练循环的核心看起来非常像入门部分中的循环。我们没有要求任何报告,因此此训练循环不会告诉我们有关模型性能的任何信息。我们需要为此添加一个评估循环。

评估循环

和之前一样,我们将使用 🤗 Evaluate 库提供的指标。我们已经看到了 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()
{'accuracy': 0.8431372549019608, 'f1': 0.8907849829351535}

同样,您的结果将略有不同,因为模型头部初始化和数据混洗存在随机性,但它们应该在同一范围内。

✏️ 动手尝试! 修改之前的训练循环,以便在 SST-2 数据集上微调您的模型。

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

以下是更改的内容

+ 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 会为您处理设备放置,因此您可以删除将模型放到设备上的行(或者,如果您愿意,可以将它们更改为使用accelerator.device而不是device)。

然后,将数据加载器、模型和优化器发送到accelerator.prepare()的那一行完成了大部分工作。这将用适当的容器包装这些对象,以确保您的分布式训练按预期进行。要做的其他更改是删除将批处理放到device上的行(同样,如果您想保留它,只需将其更改为使用accelerator.device),并用accelerator.backward(loss)替换loss.backward()

⚠️ 为了从 Cloud TPU 提供的加速中获益,我们建议您使用分词器的padding="max_length"max_length参数将您的样本填充到固定长度。

如果您想复制粘贴它来玩玩,以下是使用 🤗 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 存储库 中找到更多示例。