LLM 课程文档

一个完整的训练循环

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

一个完整的训练循环

Ask a Question Open In Colab Open In Studio Lab

现在我们将看到如何在不使用 Trainer 类的情况下实现与上一节相同的结果,我们将使用现代 PyTorch 最佳实践从头开始实现一个训练循环。同样,我们假设您已经完成了第 2 节中的数据处理。这里有一个简短的总结,涵盖了您需要的所有内容

🏗️ 从头开始训练:本节建立在之前内容的基础上。有关 PyTorch 训练循环和最佳实践的全面指导,请查阅 🤗 Transformers 训练文档自定义训练手册

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)

准备训练

在实际编写训练循环之前,我们需要定义一些对象。首先是我们用于迭代批次的 dataloaders。但在定义这些 dataloaders 之前,我们需要对 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"]

完成此操作后,我们可以轻松定义 dataloaders

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

当提供 labels 时,所有 🤗 Transformers 模型都会返回损失,我们还会得到 logits(我们的批次中每个输入有两个,因此是大小为 8 x 2 的张量)。

我们几乎可以编写训练循环了!我们只缺少两样东西:一个优化器和一个学习率调度器。由于我们正在尝试手动复制 Trainer 的行为,因此我们将使用相同的默认值。Trainer 使用的优化器是 AdamW,它与 Adam 相同,但在权重衰减正则化方面有所不同(请参阅 Ilya Loshchilov 和 Frank Hutter 的“解耦权重衰减正则化”

from torch.optim import AdamW

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

💡 现代优化技巧:为了获得更好的性能,您可以尝试

  • 带有权重衰减的 AdamWAdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
  • 8 位 Adam:使用 bitsandbytes 进行内存高效优化
  • 不同的学习率:对于大型模型,较低的学习率(1e-5 到 3e-5)通常效果更好

🚀 优化资源:在 🤗 Transformers 优化指南中了解更多关于优化器和训练策略的信息。

最后,默认使用的学习率调度器只是从最大值(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)
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)

💡 现代训练优化:为了使您的训练循环更高效,请考虑

  • 梯度裁剪:在 optimizer.step() 之前添加 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
  • 混合精度:使用 torch.cuda.amp.autocast()GradScaler 进行更快的训练
  • 梯度累积:累积多个批次的梯度以模拟更大的批次大小
  • 检查点:定期保存模型检查点,以便在中断时恢复训练

🔧 实施指南:有关这些优化的详细示例,请参阅 🤗 Transformers 高效训练指南优化器范围

您可以看到训练循环的核心与导论中的训练循环非常相似。我们没有要求任何报告,因此此训练循环不会告诉我们模型表现如何。我们需要为此添加一个评估循环。

评估循环

和前面一样,我们将使用 🤗 Evaluate 库提供的指标。我们已经看过 metric.compute() 方法,但实际上,指标可以在我们进行预测循环时使用 add_batch() 方法为我们累积批次。一旦我们累积了所有批次,我们就可以使用 metric.compute() 获得最终结果。以下是如何在评估循环中实现所有这些:

📊 评估最佳实践:有关更复杂的评估策略和指标,请浏览 🤗 Evaluate 文档综合评估手册

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 上启用分布式训练。🤗 Accelerate 自动处理分布式训练、混合精度和设备放置的复杂性。从创建训练和验证数据加载器开始,我们的手动训练循环如下所示:

Accelerate 深入探讨:在 🤗 Accelerate 文档 中了解有关分布式训练、混合精度和硬件优化的所有信息,并在 transformers 文档 中探索实际示例。

from accelerate import Accelerator
from torch.optim import AdamW
from transformers import 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)

要添加的第一行是 import 行。第二行实例化一个 Accelerator 对象,该对象将查看环境并初始化适当的分布式设置。🤗 Accelerate 会为您处理设备放置,因此您可以删除将模型放置在设备上的行(或者,如果您愿意,可以将它们更改为使用 accelerator.device 而不是 device)。

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

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

如果您想复制粘贴并尝试一下,这就是使用 🤗 Accelerate 的完整训练循环:

from accelerate import Accelerator
from torch.optim import AdamW
from transformers import 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 存储库中找到更多示例。

🌐 分布式训练:有关多 GPU 和多节点训练的全面介绍,请查阅 🤗 Transformers 分布式训练指南扩展训练手册

下一步和最佳实践

现在您已经了解了如何从头开始实现训练,以下是生产使用的一些额外考虑事项

模型评估:始终在多个指标上评估您的模型,而不仅仅是准确性。使用 🤗 Evaluate 库进行全面的评估。

超参数调优:考虑使用 Optuna 或 Ray Tune 等库进行系统的超参数优化。

模型监控:在整个训练过程中跟踪训练指标、学习曲线和验证性能。

模型共享:训练完成后,在 Hugging Face Hub 上共享您的模型,以便社区可以使用。

效率:对于大型模型,请考虑梯度检查点、参数高效微调(LoRA、AdaLoRA)或量化方法等技术。

我们对使用自定义训练循环进行微调的深入探讨到此结束。您在这里学到的技能将在您需要完全控制训练过程或想要实现超出 Trainer API 范围的自定义训练逻辑时派上用场。

本节测验

测试您对自定义训练循环和高级训练技术的理解

1. Adam 和 AdamW 优化器之间的主要区别是什么?

2. 在训练循环中,操作的正确顺序是什么?

3. 🤗 Accelerate 库主要有什么作用?

4. 为什么我们在训练循环中将批次移动到设备上?

5. 在评估之前,model.eval() 的作用是什么?

6. torch.no_grad() 在评估期间的目的是什么?

7. 在您的训练循环中使用 🤗 Accelerate 后会有哪些变化?

💡 要点:

  • 手动训练循环提供完全控制,但需要理解正确的顺序:前向 → 反向 → 优化器步进 → 调度器步进 → 梯度清零
  • 带有权重衰减的 AdamW 是推荐用于 Transformer 模型的优化器
  • 在评估期间始终使用 model.eval()torch.no_grad() 以获得正确的行为和效率
  • 🤗 Accelerate 通过最少的代码更改即可实现分布式训练
  • 设备管理(将张量移动到 GPU/CPU)对于 PyTorch 操作至关重要
  • 混合精度、梯度累积和梯度裁剪等现代技术可以显著提高训练效率
< > 在 GitHub 上更新