文章目录
完整训练
现在,我们将看到如何在不使用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进行一些后处理,以处理训练器自动为我们做的一些事情。具体来说,我们需要:
- 删除与模型不期望的值相对应的列(如sentence1和sentence2列)。
- 将列标签重命名为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,并且我们在批处理中填充到最大长度。
现在我们已经完全完成了数据预处理(对于任何ML从业者来说,这是一个令人满意但难以实现的目标),让我们转向模型。我们像上一节一样实例化它:
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])
当提供标签时,所有transformer模型都将返回损失,并且我们还获得logits(批处理中的每个输入两个,因此大小为8 x 2的张量)。
我们几乎已经准备好编写训练循环了!我们只缺少两样东西:一个优化器和一个学习率调度器。因为我们试图复制什么训练师正在做的手,我们将使用相同的默认值。训练器使用的优化器是AdamW,它与Adam相同,但对权重衰减正则化有所改变(参见Ilya Loshchilov和Frank Hutter的"解耦权重衰减正则化"):
from torch.optim 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)
结果:
1377
循环训练
最后一件事:如果我们可以使用GPU(在CPU上,训练可能需要几个小时而不是几分钟),我们将希望使用GPU。为了做到这一点,我们定义了一个设备,我们将把我们的模型和批次放在上面:
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()方法遍历预测循环时,metrics实际上可以为我们累积批。一旦我们累积了所有批次,我们就可以使用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}
同样,由于模型头部初始化和数据洗牌的随机性,您的结果将略有不同,但它们应该在相同的范围内。
训练加速
我们之前定义的训练循环在单个CPU或GPU上运行良好。但是使用加速库,只需进行一些调整,我们就可以在多个gpu或tpu上进行分布式训练。从创建训练和验证数据加载器开始,下面是我们的手动训练循环:
from torch.optim import AdamW
from transformers import 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 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)
- 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()的行中完成了大部分工作。这将把这些对象包装在适当的容器中,以确保您的分布式培训按预期工作。剩下要做的更改是删除将批处理放在设备上的行(同样,如果你想保留它,你可以将其更改为使用accelerator.device),并将loss.backward()替换为accelerator.backward(loss)。
如果你想复制粘贴它来玩,这里是完整的训练循环看起来像加速:
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
这将启动分布式培训。
如果你想在笔记本中尝试这一点(例如,用Colab上的tpu测试它),只需将代码粘贴在training_function()中并运行最后一个单元格:
from accelerate import notebook_launcher
notebook_launcher(training_function)