微调预训练模型

使用预训练模型有很多好处。预训练模型节省了你的计算开销、你的碳排放,并且让你能够使用sota模型而不需要自己从头训练。Hugging Face Transformers为你提供了上千种预训练模型,可广泛用于各种任务。当你使用一个预训练模型,你可以在任务特定数据集上训练。这就是著名的微调,一种非常厉害的训练技巧。在本篇教程中,你可以用Pytorch微调一个预训练模型。

准备一个数据集

在你微调一个预训练模型之前,下载一个数据集并将其处理成可用于训练的形式。

首先下载一个Yelp Reviews数据集:

from datasets import load_dataset

dataset = load_dataset("yelp_review_full")
dataset[100]

'''
output:
{'label': 0,
 'text': 'My expectations for McDonalds are t rarely high...}
'''

正如你所知道的那样,你需要一个分词器来处理文本以及填充、截断策略来处理可变序列长度。为了一步处理你的数据集,使用Dataset的map方法,将预处理函数应用在整个数据集上。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")


def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


tokenized_datasets = dataset.map(tokenize_function, batched=True)

你还可以使用完整数据集的一个子集来进行微调,这样可以减少时间。

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

用Trainer进行微调

Hugging Face Transformers提供了一个Trainer类用来优化模型的训练,不用人工编写自己的训练循环就能开始训练。Trainer AIP支持广泛的训练选项和特征,比如:日志、梯度积累、混合精度。

一开始要加载你的模型,并制定期望的标签数目。从Yelp Review数据集卡片得知,总共有5个标签:

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

你将看到一个“一些预训练权重没有被使用,并且一些权重被随机初始化”的提醒。不用担心,这非常正常!BERT模型的预训练头被丢弃了,取而代之的是一个随机初始化的分类头。你将会在你的序列分类任务上微调这个新模型头,将预训练模型的知识迁移过去。

训练超参数

下一步,创建一个TrainingArguments类,这个类包含了你可以调整的所有超参数,以及激活不同训练选项的标志。在本篇教程中,你可以从默认的训练超参数开始,然后进一步可以尽情实验以找到自己最佳的环境。

指定要把训练产生的checkpoint存到哪个地址:

from transformers import TrainingArguments

training_args = TrainingArguments(output_dir="test_trainer")

评价指标

Trainer不能在训练过程中自动评价模型的表现。你需要传给Trainer一个函数来计算和报告指标。Datasets库提供了一个简单的accuracy函数,可以通过load_metric()加载。

import numpy as np
from datasets import load_metric

metric = load_metric("accuracy")

在metric上调用compute可以计算模型预测的准确率。在将你的预测值放入comput之前,你需要将预测值转换为概率(记住所有的transformers模型都返回logits)。

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

如果你希望在微调过程中监控你的评价指标,在训练参数中指定evaluation_strategy参数,在每一回合最后报告评价指标

from transformers import TrainingArguments

training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")

Trainer

使用你的模型、训练参数、训练和测试数据集以及评价函数,创建一个Trainer类。

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

然后通过调用train()微调你的模型:

trainer.train()

在原生Pytorch上进行微调

Trainer包含了训练循环,并允许你使用一行代码微调模型。对于喜欢使用自己训练循环的用户,你可以在原生Pytorch上微调模型。

手工后处理tokenized_dataset

  1. 将text列删除,因为模型不能直接接收文本作为输入
tokenized_datasets = tokenized_datasets.remove_columns(["text"])
  1. 将label列重命名为labels,因为模型期望的参数名是labels
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
  1. 设置数据集的格式,使得返回张量,而不是lists
tokenized_datasets.set_format("torch")

DataLoader

为你的训练和测试数据集创建一个DataLoader,这样你可以迭代批次的数据

from torch.utils.data import DataLoader

train_dataloader = DataLoader(small_train_dataset, shuffle=True, batch_size=8)
eval_dataloader = DataLoader(small_eval_dataset, batch_size=8)

优化器和学习率衰减

创建一个优化器和学习率规划器来微调模型。优化器使用AdamW

from torch.optim import AdamW

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

创建一个默认的学习率规划器

from transformers import get_scheduler

num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)

最后,指定device使用GPU

import torch

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

现在,你就可以训练了。

训练循环

为了跟踪你的训练过程,使用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)

评价指标

和Trainer一样,在你写自己的训练循环时需要做相同的事情。但是这次,你需要将所有批次数据的数值累计起来,在最后计算指标。

metric = load_metric("accuracy")
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()