完整训练
现在我们将看到如何在不使用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)
准备训练
在实际编写训练循环之前,我们需要定义一些对象。第一个是我们将用来迭代批次的 dataloader。但在我们可以定义这些 dataloader 之前,我们需要对我们的tokenized_datasets
应用一些后处理,以处理一些Trainer
自动为我们完成的事情。具体来说,我们需要
- 删除与模型不期望的值(如
sentence1
和sentence2
列)对应的列。 - 将列
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"]
现在,我们可以轻松地定义我们的 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
{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
,并且我们在批次内部填充到最大长度,因此实际形状可能会有所不同。
现在,我们已经完成了数据预处理(对于任何机器学习从业人员来说,这是一个令人满意但难以捉摸的目标),让我们转向模型。我们实例化它的方式与上一节完全相同
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 的调度器。为了正确定义它,我们需要知道我们将采取的训练步骤数,即我们想要运行的 epoch 数乘以训练批次数(即训练 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
训练循环
最后,我们希望使用 GPU(如果我们有权访问 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 上启用分布式训练。从创建训练和验证 dataloader 开始,以下是手动训练循环的样子
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
)。
然后,主要工作量是在将 dataloaders、模型和优化器发送到accelerator.prepare()
的那一行中完成的。这将用适当的容器包装这些对象,以确保您的分布式训练按预期进行。要做的其余更改是删除将批处理放在device
上的代码行(同样,如果您想保留它,只需将其更改为使用accelerator.device
)并将loss.backward()
替换为accelerator.backward(loss)
。
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
这将启动分布式训练。
如果您想在笔记本(例如在 Colab 上使用 TPU 测试它)中试用它,只需将代码粘贴到training_function()
中,然后运行最后一个单元格:
from accelerate import notebook_launcher
notebook_launcher(training_function)
您可以在🤗 Accelerate 代码库中找到更多示例。