BART预训练模型实战入门：从零构建文本生成任务的完整流程

1次阅读

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

BART（Bidirectional and Auto-Regressive Transformers）是一种基于 Transformer 的 seq2seq（序列到序列）预训练模型，它结合了双向编码器和自回归解码器的优势。BART 在文本生成任务中表现出色，如文本摘要、对话生成等，因为它能有效建模输入文本的上下文信息并生成流畅的输出。相比其他生成模型，BART 在微调阶段对数据量的需求较低，适合资源有限的开发者快速实现生产级应用。

使用 HuggingFace 的 transformers 库加载 BART-large 模型时，显存占用可能高达 3GB，这对于显存有限的 GPU（如 16GB）来说是一个挑战。我们可以通过梯度检查点（Gradient Checkpointing）技术减少显存占用，同时保持模型性能。

from transformers import BartForConditionalGeneration
import torch

model = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
model.config.gradient_checkpointing = True  # 启用梯度检查点
model.train()

在微调 BART 时，学习率设置不当可能导致训练不稳定，表现为 loss 剧烈波动。建议使用带 Warmup 的 AdamW 优化器，逐步增加学习率以避免初始震荡。

from transformers import AdamW, get_linear_schedule_with_warmup

optimizer = AdamW(model.parameters(), lr=5e-5, eps=1e-8)
total_steps = len(train_dataloader) * epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=int(total_steps * 0.1),  # Warmup 10% 的总步数
    num_training_steps=total_steps
)

BART 在生成文本时可能出现重复或短句问题，可以通过调整 Beam Search 参数和引入 N -gram 惩罚来缓解。

generated = model.generate(
    input_ids, 
    max_length=50, 
    num_beams=5,  # Beam 宽度设为 5
    no_repeat_ngram_size=2,  # 禁止 2 -gram 重复
    early_stopping=True
)

在 16GB 显存的 GPU 上，启用梯度检查点后，BART-large 的最大 batch size 可提升至 8（输入长度 128）。下表展示了不同输入长度下的显存占用情况：

输入长度	Batch Size	显存占用（GB）
64	16	12.3
128	8	14.1
256	4	15.8

测试不同输入长度下的推理延迟（P50 和 P90 百分位数）：

输入长度 64：P50=120ms，P90=150ms
输入长度 128：P50=210ms，P90=280ms
输入长度 256：P50=450ms，P90=580ms

使用混合精度训练（FP16）时，可能出现 Loss 为 NaN 的情况。解决方法：

torch.cuda.amp.GradScaler().scale(loss).backward()  # 使用 GradScaler

BART 的默认 Tokenizer 对中文分字可能不够友好，建议对输入文本手动分字或使用中文专用 Tokenizer：

from transformers import BartTokenizer

tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
text = "这是一个例子"
tokens = tokenizer.tokenize(text)  # 输出可能不如预期