BERT预训练模型实现指南：从零搭建到性能调优

1次阅读

没有评论

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

对于刚接触 BERT 预训练的开发者，通常会遇到以下三类典型问题：

多 GPU 并行训练配置复杂：数据并行与模型并行的选择、梯度同步策略、DistributedDataParallel 的初始化方式等细节容易出错
长文本处理效率低下：直接截断会导致信息丢失，而使用稀疏注意力或分块机制又需要修改模型结构
自定义词典接入困难：原生的 WordPiece 分词器对中文支持有限，重新训练 tokenizer 需要处理字符覆盖率和词汇表平衡问题

关于实现方式的选择，Hugging Face Transformers 库适合：

快速验证想法
需要兼容多种预训练模型
工业级部署场景

手动 PyTorch 实现则更适合：

研究模型改进
特殊硬件适配
教学演示场景

import torch
import math

class BERTEmbedding(torch.nn.Module):
    def __init__(self, vocab_size, hidden_size, max_position_embeddings):
        super().__init__()
        # 词向量、位置向量、token 类型向量三部分叠加
        self.word_embeddings = torch.nn.Embedding(vocab_size, hidden_size)
        self.position_embeddings = torch.nn.Embedding(max_position_embeddings, hidden_size)
        self.token_type_embeddings = torch.nn.Embedding(2, hidden_size)  # 通常只需区分两个句子

        # LayerNorm 和 Dropout 是 BERT 稳定训练的关键
        self.LayerNorm = torch.nn.LayerNorm(hidden_size, eps=1e-12)
        self.dropout = torch.nn.Dropout(0.1)

    def forward(self, input_ids, token_type_ids=None, position_ids=None):
        seq_length = input_ids.size(1)

        if position_ids is None:
            # 自动生成位置 ID [0,1,2,...,seq_len-1]
            position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        # 三部分向量相加
        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = words_embeddings + position_embeddings + token_type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

负采样策略建议采用：

80% 的概率替换为 [MASK] 标记
10% 的概率替换为随机词
10% 的概率保持原词不变

损失函数计算示例：

def mlm_loss(hidden_states, labels, vocab_size):
    # hidden_states: [batch_size, seq_len, hidden_size]
    # labels: [batch_size, seq_len] 其中未被 mask 的位置为 -100
    mlm_dense = torch.nn.Linear(hidden_size, vocab_size)
    loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100)

    logits = mlm_dense(hidden_states)  # [batch_size, seq_len, vocab_size]
    loss = loss_fct(logits.view(-1, vocab_size), labels.view(-1))
    return loss

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()

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

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

def collate_fn(batch):
    max_len = max(len(x['input_ids']) for x in batch)

    # 预分配张量避免多次扩容
    input_ids = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
    attention_mask = torch.zeros(len(batch), max_len, dtype=torch.long)

    for i, item in enumerate(batch):
        length = len(item['input_ids'])
        input_ids[i, :length] = torch.tensor(item['input_ids'])
        attention_mask[i, :length] = 1

    return {'input_ids': input_ids, 'attention_mask': attention_mask}

检查梯度爆炸：添加torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
验证输入数据：是否存在异常值或未归一化的特征
降低初始学习率：尝试从 5e- 6 开始逐步上调

from transformers import get_linear_schedule_with_warmup

# 典型配置：10% 的 step 用于 warmup
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=int(0.1 * total_steps),
    num_training_steps=total_steps
)