BERT模型实战：从零开始用PyTorch复现Transformer Encoder模块

BERT模型实战：从零开始用PyTorch复现Transformer Encoder模块

如果你已经对Transformer架构有所耳闻，甚至用过Hugging Face的transformers库轻松调用BERT，但内心深处总有个声音在问：“这玩意儿到底是怎么转起来的？” 那么，这篇文章就是为你准备的。我们不再满足于当一个API调用者，而是要卷起袖子，用PyTorch亲手把BERT的核心——Transformer Encoder——从一个个张量运算开始搭建起来。这个过程，就像在组装一台精密的钟表，你会亲眼看到位置编码如何注入序列的顺序感，多头注意力机制如何让模型同时聚焦于文本的不同侧面，以及层归一化和前馈网络如何协同工作。最终，你得到的不仅是一段能运行的代码，更是一种对模型运作机理的深刻直觉，这种直觉在模型调试、改进乃至创新时，是无价的。

1. 环境准备与项目初始化

动手之前，确保你的工作环境是整洁且可复现的。我强烈建议使用Conda来管理Python环境，它能有效避免不同项目间的依赖冲突。

首先，创建一个新的Conda环境并安装核心依赖。这里我们选择PyTorch 1.13.1和CUDA 11.6（如果你的机器支持GPU），这能保证大多数算子的稳定性和兼容性。

conda create -n bert_from_scratch python=3.9 -y
conda activate bert_from_scratch
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install numpy matplotlib tqdm

接下来，初始化你的项目结构。一个清晰的结构能让代码逻辑一目了然，也便于后续扩展。我习惯的目录结构如下：

bert_from_scratch/
├── config.py        # 模型参数配置
├── model/
│   ├── __init__.py
│   ├── embedding.py # 词嵌入 & 位置编码
│   ├── attention.py # 注意力机制
│   ├── encoder.py   # Encoder层及堆叠
│   └── bert.py      # BERT模型组装
├── utils.py         # 工具函数（如GELU激活）
├── train_demo.py    # 简易训练/推理演示
└── README.md

在config.py中，我们先定义一组基础的模型配置参数，这模仿了BERT-Base的设定，也方便后续调整。

# config.py
class BertConfig:
    def __init__(self):
        self.vocab_size = 30522  # BERT原始词表大小
        self.hidden_size = 768   # 隐藏层维度
        self.num_hidden_layers = 12  # Transformer Encoder层数
        self.num_attention_heads = 12  # 注意力头数
        self.intermediate_size = 3072  # 前馈网络中间层维度
        self.hidden_act = "gelu"       # 激活函数
        self.hidden_dropout_prob = 0.1  # 隐藏层dropout率
        self.attention_probs_dropout_prob = 0.1  # 注意力概率dropout率
        self.max_position_embeddings = 512  # 最大序列长度
        self.type_vocab_size = 2  # 句子类型（如Segment A/B）数量
        self.initializer_range = 0.02  # 参数初始化范围
        self.layer_norm_eps = 1e-12    # LayerNorm的epsilon

提示：将配置参数集中管理是一个好习惯。当你想尝试不同规模的模型（如 TinyBERT）时，只需在此处创建新的配置类，而无需深入修改模型代码。

2. 构建嵌入层：让模型“读懂”输入

BERT的输入并非原始文本，而是三种信息的融合：词元（Token）本身、它在序列中的位置、以及它所属的句子片段。我们的嵌入层需要同时处理这三者。

2.1 词元嵌入与句子类型嵌入

词元嵌入将每个词（或子词）映射到一个高维向量。句子类型嵌入则用于区分句子对中的第一句和第二句，这对下一句预测（NSP）任务至关重要。实现上，它们都是简单的nn.Embedding层。

# model/embedding.py
import torch
import torch.nn as nn
import math

class BertEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

        # 注册一个缓冲区，用于存储位置ID，避免每次前向传播都重新生成
        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))

2.2 位置编码：注入序列顺序的灵魂

与原始Transformer使用正弦余弦函数不同，BERT采用了更简单的可学习位置嵌入。这意味着模型会在训练过程中自己学会什么样的位置该对应什么样的向量。虽然少了些理论上的归纳偏置，但实践表明其效果足够好，且更易实现。

前向传播的逻辑就是将三种嵌入相加，然后进行层归一化和Dropout。

    def forward(self, input_ids, token_type_ids=None):
        """
        Args:
            input_ids: [batch_size, seq_len]
            token_type_ids: [batch_size, seq_len]， 默认为0（表示第一句）
        """
        seq_length = input_ids.size(1)
        position_ids = self.position_ids[:, :seq_length]

        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

这里有一个细节：我们使用register_buffer将position_ids注册为模块的缓冲区。这意味着它会随模型一起保存和加载，但不会被视为可训练参数。这比在每次前向传播时用torch.arange新建张量更高效。

3. 核心引擎：自注意力机制详解

自注意力机制是Transformer的灵魂，它允许序列中的每个位置直接关注到序列的所有其他位置，从而捕获长距离依赖。BERT使用的是“缩放点积注意力”。

3.1 缩放点积注意力

其公式为：Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V 其中，Q（查询）、K（键）、V（值）均由输入线性变换而来。除以sqrt(d_k)是为了防止点积结果过大导致softmax梯度消失。

# model/attention.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout_prob=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self, query, key, value, attention_mask=None):
        """
        Args:
            query: [batch_size, num_heads, seq_len, head_dim]
            key: [batch_size, num_heads, seq_len, head_dim]
            value: [batch_size, num_heads, seq_len, head_dim]
            attention_mask: [batch_size, 1, 1, seq_len] 或 [batch_size, 1, seq_len, seq_len]， 掩码处为0或极小负值
        """
        d_k = query.size(-1)
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

        if attention_mask is not None:
            # 将掩码中为0的位置替换为一个非常大的负数，使得softmax后概率接近0
            scores = scores + attention_mask

        attn_probs = F.softmax(scores, dim=-1)
        attn_probs = self.dropout(attn_probs)
        context = torch.matmul(attn_probs, value)
        return context, attn_probs

3.2 多头注意力：并行化的特征提取

单头注意力可能只关注到一种模式的关系。多头注意力将hidden_size维度的信息拆分到多个“头”上，让每个头在不同的子空间里学习不同的关注模式，最后再将结果合并。这极大地增强了模型的表征能力。

class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        assert self.head_dim * self.num_heads == config.hidden_size, "hidden_size必须能被num_attention_heads整除"

        self.query = nn.Linear(config.hidden_size, config.hidden_size)
        self.key = nn.Linear(config.hidden_size, config.hidden_size)
        self.value = nn.Linear(config.hidden_size, config.hidden_size)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.attention = ScaledDotProductAttention(config.attention_probs_dropout_prob)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def transpose_for_scores(self, x):
        # 将形状从 [batch_size, seq_len, hidden_size] 转换为 [batch_size, num_heads, seq_len, head_dim]
        new_shape = x.size()[:-1] + (self.num_heads, self.head_dim)
        x = x.view(*new_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states, attention_mask=None):
        batch_size, seq_len, _ = hidden_states.shape

        # 线性投影
        mixed_query = self.query(hidden_states)
        mixed_key = self.key(hidden_states)
        mixed_value = self.value(hidden_states)

        # 转置以准备多头计算
        query_layer = self.transpose_for_scores(mixed_query)
        key_layer = self.transpose_for_scores(mixed_key)
        value_layer = self.transpose_for_scores(mixed_value)

        # 计算缩放点积注意力
        context_layer, attn_probs = self.attention(query_layer, key_layer, value_layer, attention_mask)

        # 将多头输出合并回 [batch_size, seq_len, hidden_size]
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_shape = context_layer.size()[:-2] + (self.num_heads * self.head_dim,)
        context_layer = context_layer.view(*new_context_shape)

        # 最终投影和残差连接
        output = self.dense(context_layer)
        output = self.dropout(output)
        output = self.LayerNorm(output + hidden_states)  # 残差连接后LayerNorm
        return output, attn_probs

这里的关键操作是transpose_for_scores，它通过view和permute操作巧妙地实现了张量的重塑，使得后续的注意力计算可以并行地在所有头上进行。残差连接（output + hidden_states）和层归一化是稳定深层网络训练的关键技术，它们有助于缓解梯度消失问题。

4. 前馈网络与编码器层组装

注意力机制负责聚合信息，而前馈网络（FFN）则负责对每个位置的特征进行非线性变换和升维，为模型增加表达能力。

4.1 前馈网络的结构

BERT的前馈网络是一个两层MLP，中间层的维度（intermediate_size）通常是隐藏层维度（hidden_size）的4倍，使用GELU激活函数。

# utils.py
def gelu(x):
    """Gaussian Error Linear Unit激活函数。
    这是BERT中使用的近似计算，与原始GELU略有不同但计算更快。
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

# model/encoder.py
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.dense2 = nn.Linear(config.intermediate_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

    def forward(self, hidden_states):
        intermediate = self.dense1(hidden_states)
        intermediate = gelu(intermediate)  # 使用自定义的GELU
        intermediate = self.dense2(intermediate)
        intermediate = self.dropout(intermediate)
        output = self.LayerNorm(intermediate + hidden_states)  # 残差连接
        return output

4.2 组装完整的Transformer编码器层

一个标准的Transformer编码器层由多头自注意力子层和前馈网络子层顺序连接而成。每个子层都包裹着残差连接和层归一化。

class TransformerEncoderLayer(nn.Module):
    """一个完整的Transformer Encoder层。"""
    def __init__(self, config):
        super().__init__()
        self.attention = MultiHeadAttention(config)
        self.feed_forward = FeedForward(config)

    def forward(self, hidden_states, attention_mask=None):
        # 自注意力子层
        attention_output, attn_probs = self.attention(hidden_states, attention_mask)
        # 前馈网络子层
        encoder_output = self.feed_forward(attention_output)
        return encoder_output, attn_probs

4.3 堆叠编码器层以构建深度模型

BERT的强大能力来源于其深度。通过堆叠多个相同的编码器层，模型可以构建出越来越抽象和复杂的语言表示。

class TransformerEncoder(nn.Module):
    """堆叠N个TransformerEncoderLayer。"""
    def __init__(self, config):
        super().__init__()
        self.layer = nn.ModuleList([TransformerEncoderLayer(config) for _ in range(config.num_hidden_layers)])

    def forward(self, hidden_states, attention_mask=None, output_attentions=False):
        all_attentions = () if output_attentions else None
        for layer_module in self.layer:
            hidden_states, attn_probs = layer_module(hidden_states, attention_mask)
            if output_attentions:
                all_attentions = all_attentions + (attn_probs,)
        if output_attentions:
            return hidden_states, all_attentions
        return hidden_states

output_attentions参数是一个实用功能，它允许我们在前向传播中返回每一层的注意力权重。这对于模型可解释性分析（可视化模型关注了哪些词）非常有帮助。

5. 整合与对比：从零实现与HuggingFace的异同

现在，我们将所有部件组装成完整的BERT模型，并探讨我们的实现与主流库（如Hugging Face transformers）之间的关键差异。

5.1 组装BERT模型

除了编码器，BERT模型还需要一个池化器（Pooler）来提取句子级别的表示，主要用于下一句预测任务。

# model/bert.py
import torch.nn as nn
from .embedding import BertEmbeddings
from .encoder import TransformerEncoder

class BertModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embeddings = BertEmbeddings(config)
        self.encoder = TransformerEncoder(config)
        self.pooler = nn.Sequential(
            nn.Linear(config.hidden_size, config.hidden_size),
            nn.Tanh()
        )

    def forward(self, input_ids, attention_mask=None, token_type_ids=None, output_attentions=False):
        # 1. 创建注意力掩码（如果需要）
        if attention_mask is not None:
            # 将[1, 1, 1, 0, 0]转换为适合加到注意力分数上的形式
            extended_attention_mask = attention_mask[:, None, None, :]
            extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
        else:
            extended_attention_mask = None

        # 2. 嵌入层
        embedding_output = self.embeddings(input_ids, token_type_ids)

        # 3. 编码器层
        encoder_outputs = self.encoder(embedding_output, extended_attention_mask, output_attentions)

        if output_attentions:
            sequence_output, all_attentions = encoder_outputs
        else:
            sequence_output = encoder_outputs
            all_attentions = None

        # 4. 池化输出（取[CLS]标记对应的隐藏状态）
        pooled_output = self.pooler(sequence_output[:, 0])

        return sequence_output, pooled_output, all_attentions

5.2 关键实现细节对比

我们的实现力求清晰易懂，而工业级库如transformers则在此基础上做了大量优化。了解这些差异有助于你理解生产级代码的考量。

注意：我们的实现专注于教学和原理展示，因此在性能优化、边界条件处理和生态集成上做了简化。在实际项目中使用时，应优先考虑成熟的库，但理解其底层实现能让你在遇到问题时游刃有余。

5.3 简易前向传播验证

让我们写一小段代码来验证模型是否能正确运行，并观察输出的形状。

# train_demo.py (部分)
from config import BertConfig
from model.bert import BertModel

if __name__ == "__main__":
    config = BertConfig()
    model = BertModel(config)
    model.eval()  # 切换到评估模式

    # 模拟一个batch的输入：batch_size=2, seq_len=16
    input_ids = torch.randint(0, config.vocab_size, (2, 16))
    attention_mask = torch.ones(2, 16)  # 假设所有位置都是有效的
    attention_mask[0, 10:] = 0  # 第一个序列的后6个位置是padding

    with torch.no_grad():
        sequence_output, pooled_output, attentions = model(input_ids, attention_mask, output_attentions=True)

    print(f"输入形状: {input_ids.shape}")
    print(f"序列输出形状: {sequence_output.shape}")  # 应为 [2, 16, 768]
    print(f"池化输出形状: {pooled_output.shape}")    # 应为 [2, 768]
    print(f"注意力权重元组长度: {len(attentions)}")   # 应为 12 (层数)
    print(f"第一层注意力权重形状: {attentions[0].shape}") # 应为 [2, 12, 16, 16] (batch, heads, seq, seq)

运行这段代码，如果一切正常，你应该能看到符合预期的张量形状。这表明我们的模型在结构上是正确的。

6. 扩展思考与实战建议

构建出基础模型只是第一步。要让它在实际任务中发挥作用，或者进行更深入的探索，这里有几个方向和技巧。

从零预训练是否可行？ 理论上可以，但成本极高。BERT-Base在数十GB文本上训练了数天甚至数周，需要强大的算力。对于个人开发者，更现实的路径是：

1. 加载预训练权重进行微调：使用Hugging Face库加载官方或社区预训练好的权重，在你的特定任务（如文本分类、问答）上进行微调。这是最主流、最有效的方法。
2. 知识蒸馏：用大型教师模型（如BERT-Large）来指导我们这个小模型训练，可以在损失少量性能的情况下大幅减少模型体积和推理时间。
3. 继续预训练：在特定领域语料（如医学、法律文献）上对通用BERT进行额外的预训练，使其适应领域语言特点。

调试与可视化技巧 当模型输出不如预期时，可以尝试以下方法：

• 检查中间变量：在前向传播的关键步骤（如嵌入相加后、注意力softmax前）打印张量的均值和标准差，观察是否有异常值（如NaN或Inf）。
• 可视化注意力权重：将output_attentions设为True，获取注意力矩阵并可视化。这能直观地看到模型在处理一个句子时，每个词更关注哪些其他词。例如，对于句子“The animal didn't cross the street because it was too tired”，你可以观察it这个词的注意力是否正确地聚焦于animal。

  # 简易注意力热图绘制示例
  import matplotlib.pyplot as plt
  import seaborn as sns
  sentence = ["[CLS]", "the", "animal", "didn", "##t", "cross", "the", "street", "because", "it", "was", "too", "tired", "[SEP]"]
  # 假设attn是某一层某一头的注意力矩阵 [seq_len, seq_len]
  plt.figure(figsize=(10,8))
  sns.heatmap(attn[0].detach().numpy(), xticklabels=sentence, yticklabels=sentence, cmap="YlGnBu")
  plt.title("Attention Heatmap")
  plt.show()
  

性能优化小贴士

• 使用torch.jit.script或torch.compile：对于定义好的模型，可以尝试PyTorch的即时编译功能，这可能会提升推理速度。
• 注意力计算优化：对于极长序列，标准的注意力计算复杂度是O(n²)，会成为瓶颈。可以关注如FlashAttention等更高效的算法实现，它们能更好地利用GPU内存层次结构。
• 精度混合训练：使用torch.cuda.amp进行自动混合精度训练，可以在几乎不影响精度的情况下减少显存占用并加快训练速度。

亲手实现一遍之后，再回头去看Hugging Face的源码，你会发现那些看似复杂的代码变得亲切了许多。你不再是一个黑盒的使用者，而是一个知其所以然的构建者。下次当你的模型表现不佳时，你可能会更有底气地去检查注意力机制是否正常，或者嵌入层是否得到了合理的初始化。这种从底层积累起来的理解，正是区分普通开发者和资深专家的关键所在。