一个完整的训练
现在,我们将了解如何在不使用 Trainer
类的情况下实现与上一节相同的结果。同样,我们假设你已经完成了第 2 节中的数据处理。下面对第 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)
训练前的准备
在正式开始编写我们的训练循环之前,我们需要定义一些对象。首先是我们将用于迭代 batch 的数据加载器。但在定义这些数据加载器之前,我们需要对我们的 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"]
至此,我们可以轻松定义数据加载器:
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
)
为了快速检验数据处理中没有错误,我们可以这样检验其中的一个 batch:
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
,并且模型会将句子填充到 batch
中的最大长度。
现在我们已经完全完成了数据预处理(对于任何 ML 从业者来说都是一个令人满意但难以实现的目标),让我们将注意力转向模型。我们会像在上一节中所做的那样实例化它:
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])
当我们输入 labels
时,🤗 Transformers 模型都将返回这个 batch
的 loss
,我们还得到了 logits
( batch
中的每个输入有两个输出,所以张量大小为 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 的线性衰减。为了定义它,我们需要知道我们训练的次数,即所有数据训练的次数(epochs)乘以的 batch 的数量(即我们训练数据加载器的长度)。 Trainer
默认情况下使用三个 epochs
,因此我们定义训练过程如下:
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
,它在 GPU 可用的情况下指向 GPU,最后我们将把我们的模型和 batch
放在 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()
方法进行预测循环时,实际上该指标可以为我们累积所有 batch
的结果。一旦我们累积了所有 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
。第二行实例化一个 Accelerator
对象 它将查看环境并初始化适当的分布式设置。🤗 Accelerate 为你处理数据在设备间的数据传递,因此你可以删除将模型放在设备上的那行代码(或者,如果你愿意,可使用 accelerator.device
代替 device
)。
然后大部分工作会在将数据加载器、模型和优化器发送到的 accelerator.prepare()
中完成。这将会把这些对象包装在适当的容器中,以确保你的分布式训练按预期工作。要进行的其余更改是删除将 batch
放在 device
的那行代码(同样,如果你想保留它,你可以将其更改为使用 accelerator.device
) 并将 loss.backward()
替换为 accelerator.backward(loss)
。
⚠️ 为了使云端 TPU 提供的加速中发挥最大的效益,我们建议使用 tokenizer 的 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 repo 找到更多的示例。
< > Update on GitHub