Transformers源码解析 - 文档篇 - 模型训练

微调预训练模型

大模型的参数量非常大，个人用户训练整个模型是不实际的，实践上一般都是微调预训练过的模型。

数据准备

加载数据集：

from datasets import load_dataset

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

加载并使用tokenizer。dataset的map方法可以一次性处理整个数据集，适用于不需要做数据增强的情况，跟上一节讲到的dataset.set_transform()形成了对比。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("google-bert/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))

使用transformers库进行训练

transformers提供了一个专为训练 Transformers 模型而优化的 Trainer 类，使您无需手动编写自己的训练循环步骤而更轻松地开始训练模型。Trainer API 支持各种训练选项和功能，如日志记录、梯度累积和混合精度。

初始化模型

首先是初始化模型，一般只需要修改num_labels，调整分类头的输出单元数量：

from transformers import AutoModelForSequenceClassification

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

预训练模型的分类头将被替换为全新初始化的分类头。

配置训练超参数

通过TrainingArguments类来构造一组训练超参数，都是有默认值的。

from transformers import TrainingArguments

training_args = TrainingArguments(output_dir="test_trainer")
# 每个epoch结束进行一次评估，关于评估环节后面的内容会讲到
training_args = TrainingArguments(output_dir="test_trainer", eval_strategy="epoch")

评估

如果我们需要在训练过程中评估模型的性能，那么需要自己准备指标计算函数。
首先是构造一个metric，evaluate库中提供了一系列指标，我们这里选用一个正确率"accuracy"指标：

import numpy as np
import evaluate

metric = evaluate.load("accuracy")

光有指标还不够，我们要写好在接收到模型的输出之后，如何将输出传入到metric：

# 注意：所有的transformers库里的模型推理时都返回logits对
def compute_metrics(eval_pred):
    # eval_pred就是模型推理时的输出
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    # 调用我们前面构造的metric的compute接口
    return metric.compute(predictions=predictions, references=labels)

训练器

构造一个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,
)

训练，启动！

trainer.train()

使用keras训练tensorflow模型一节略过。

使用原生pytorch进行训练

接下来我们再来看看，如果使用原生pytorch进行训练，需要做哪些改变。

数据集

tokenized_datasets需要做一些改造：

# 移除 text 列，因为模型不接受原始文本作为输入
tokenized_datasets = tokenized_datasets.remove_columns(["text"])
# 将 label 列重命名为 labels，因为模型期望参数的名称为 labels
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
# 设置数据集的格式以返回 PyTorch 张量而不是lists
tokenized_datasets.set_format("torch")

改造完成，跟之前一样构造一个小数据集：

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

数据加载器

dataset对象可以直接用于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)

初始化模型

这一步跟用transforms训练是一样的：

from transformers import AutoModelForSequenceClassification

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

优化器

原生pytorch需要自己构造优化器和学习率调节器，transformers应该是都帮我们代劳了。

from torch.optim import AdamW
from transformers import get_scheduler

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

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
)

记得指定设备：

import torch

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

训练循环

显示进度、加载数据、前向传播得到loss、反向传播、各个优化器进行step操作：

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)

评估

原生pytorch同样可以使用evaluate库来进行评估，区别在于这次要手动调用metric.add_batch来收集模型的输出，最后调用compute来计算指标：

import evaluate

metric = evaluate.load("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()

通过已有的脚本训练模型（不重要）

transformers仓库提供了一些训练模型的脚本，不过现在似乎不怎么维护了。通过脚本训练模型需要把源代码拉下来。

git clone https://github.com/huggingface/transformers
cd transformers
pip install .

切换transformers的版本到脚本所要求的版本，并安装好依赖：

git checkout tags/v3.5.1
pip install -r requirements.txt

示例脚本通过Datasets库下载并预处理数据集。然后，脚本通过Trainer使用支持摘要任务的架构对数据集进行微调。以下示例展示了如何在CNN/DailyMail数据集上微调T5-small。由于T5模型的训练方式，它需要一个额外的source_prefix参数来让T5知道这是一个摘要任务。

python examples/pytorch/summarization/run_summarization.py \
    --model_name_or_path google-t5/t5-small \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate

Trainer支持分布式训练和混合精度，通过--fp16启用混合精度训练，通过--nproc_per_node 8设置使用GPU的数量。

torchrun \
    --nproc_per_node 8 pytorch/summarization/run_summarization.py \
    --fp16 \
    --model_name_or_path google-t5/t5-small \
    --do_train \
    --do_eval \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir /tmp/tst-summarization \
    --per_device_train_batch_size=4 \
    --per_device_eval_batch_size=4 \
    --overwrite_output_dir \
    --predict_with_generate

TPU的训练方式略过。

基于Accelerate运行脚本

Accelerate 是一个仅支持 PyTorch 的库，它提供了一种统一的方法来在不同类型的设置（仅 CPU、多个 GPU、多个TPU）上训练模型，同时保持对 PyTorch 训练循环的完全可见性。
accelerate还处于快速迭代阶段，因此建议使用源码安装：

pip install git+https://github.com/huggingface/accelerate

Accelerate支持的脚本需要在文件夹中有一个task_no_trainer.py文件，因此你需要使用run_summarization_no_trainer.py而不是run_summarization.py。首先运行以下命令以创建并保存配置文件：

accelerate config
# 检测您的设置以确保配置正确：
accelerate test

然后就可以开始训练模型了：

accelerate launch run_summarization_no_trainer.py \
    --model_name_or_path google-t5/t5-small \
    --dataset_name cnn_dailymail \
    --dataset_config "3.0.0" \
    --source_prefix "summarize: " \
    --output_dir ~/tmp/tst-summarization

使用自定义数据集

训练脚本支持使用自定义数据集，比如摘要任务就支持用户自己选择CSV或json line文件。用户需要根据任务的需求，在原有训练脚本命令的基础上增加一些额外的参数。比如：
用--train_file xx指定训练文件的路径，--validation_file指定验证文件的路径；
用--text_column指定需要进行摘要的列，此即为输入；
用--summary_column指定目标输出的文本，此即为gt。

测试脚本

这是一个常用的功能，运行完整数据集之前，先在小数据集上进行测试是合理的。通过--max_train_samples、--max_eval_samples对最大的训练和验证样本数进行限制，快速跑通全流程。--max_predict_samples只有部分脚本才有这个参数，略过。

从checkpoint恢复训练

这也是一个常用的功能，这将确保在训练中断时，您可以从之前停止的地方继续进行，而无需重新开始。
方案一是删掉--overwrite_output_dir后重跑原来的训练脚本，会自动从--output_dir中读取checkpoint进行resume。切记要把overwrite参数给删了，不然就莫得了。
方案二是添加一个额外的--resume_from_checkpoint xx参数来指定一个目录，从中读取checkpoint进行resume。

分享模型

脚本还配备了将训练好的模型上传到modelhub的选项。训练之前要先配置好huggingface hub：

huggingface-cli login

添加--push_to_hub来启用模型上传功能，通过--push_to_hub_model_id来指定要上传到云端的哪一个仓库里面，它会将--output-dir中的内容上传。

分布式训练

本节主要介绍如何使用accelerate库进行分布式训练。

首先是导入并创建Accelerator对象。Accelerator将自动检测您的分布式设置类型，并初始化所有必要的训练组件。您不需要显式地将模型放在设备上。

from accelerate import Accelerator

accelerator = Accelerator()

下一步是将所有相关的训练对象传递给prepare方法，包括训练和评估DataLoader、一个模型和一个优化器：

train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
    train_dataloader, eval_dataloader, model, optimizer
)

最后是反向传播的时候用accelerate的backward方法替换训练循环中的典型loss.backward():

for epoch in range(num_epochs):
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

下图清晰展示了使用accelerate之后，代码上所做的更改。新增四处代码，删去三处代码，即可用上accelerate了，看起来是非常的便捷。

开启训练！

# 创建和保存配置文件
accelerate config
# 启动训练
accelerate launch train.py