BERT预训练模型在SQuADv2.0数据集上的训练实战：从数据准备到性能调优

1次阅读

共计 1982 个字符，预计需要花费 5 分钟才能阅读完成。

SQuADv2.0 相较于 v1.1 版本主要增加了不可回答问题（unanswerable questions）的挑战。这要求模型不仅要找到答案的位置，还要判断问题是否可回答。此外，SQuADv2.0 中的段落通常较长，这对 BERT 模型的 512 个 token 长度限制提出了挑战。

对于初学者来说，BERT-base 是一个更好的起点：

BERT-base 有 12 层，110M 参数，在大多数消费级 GPU（如 RTX 3090）上可以运行
BERT-large 有 24 层，340M 参数，需要更大的显存（通常需要 A100 级别的 GPU）
在推理速度上，BERT-base 比 BERT-large 快约 2 - 3 倍

使用 HuggingFace 的 DatasetDict 可以方便地加载 SQuADv2.0 数据集：

from datasets import load_dataset
squad_dataset = load_dataset("squad_v2")

对于每个可回答问题，我们需要找到答案在 tokenized 文本中的 start 和 end 位置：

首先对问题和上下文进行 tokenize
然后找到原始文本中答案的字符位置
将这些字符位置映射到 tokenized 后的位置

对于不可回答问题：

将 answer_start 和 answer_end 都设为 0
在训练时，模型需要学习输出 [CLS] 位置的 start 和 end

使用 DataCollatorWithPadding 可以自动处理不同长度的样本：

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

训练开始时逐渐增大学习率，有助于稳定训练：

from transformers import AdamW, get_linear_schedule_with_warmup

optimizer = AdamW(model.parameters(), lr=5e-5)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=500,
    num_training_steps=len(train_dataloader)*epochs
)

当 GPU 显存不足时，可以通过梯度累积来等效增大 batch size：

for step, batch in enumerate(train_dataloader):
    outputs = model(**batch)
    loss = outputs.loss
    loss = loss / gradient_accumulation_steps
    loss.backward()

    if step % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

在 RTX 3090 上（24GB 显存）：

BERT-base, batch size=8: 约 10GB
BERT-base, batch size=16: 约 18GB
BERT-large, batch size=2: 约 15GB

使用 FP16 可以显著减少显存占用并提高训练速度：

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

with autocast():
    outputs = model(**batch)
    loss = outputs.loss

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

对于超过 512 token 的段落：