用PythonPyTorch手搓BERT核心5分钟掌握双向Transformer精髓BERT模型自2018年问世以来已成为自然语言处理领域的基石技术。但很多开发者发现仅通过论文和理论讲解很难真正理解其双向编码的魔力。本文将带您用不到50行Python代码实现一个微型BERT的核心功能通过可运行的代码揭示Masked Language Model(MLM)的训练奥秘。1. 准备工作理解简化版BERT的设计在开始编码前我们需要明确这个简化版BERT的定位。完整BERT-base模型有1.1亿参数而我们实现的微型版本将保留以下核心特征双向Transformer编码器使用自注意力机制同时处理左右上下文Masked Language Model通过预测被遮盖的单词学习上下文表征位置编码保留原始Transformer的位置信息处理能力import torch import torch.nn as nn import math # 超参数设置 VOCAB_SIZE 10000 # 简化词表大小 EMBED_DIM 128 # 嵌入维度 N_LAYERS 2 # Transformer层数 N_HEADS 4 # 注意力头数 MAX_LEN 64 # 最大序列长度2. 构建核心组件从嵌入层到Transformer真正的BERT使用WordPiece分词我们简化使用普通词嵌入。关键是要实现位置编码让模型理解单词顺序class PositionalEncoding(nn.Module): def __init__(self, d_model, max_lenMAX_LEN): super().__init__() position torch.arange(max_len).unsqueeze(1) div_term torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe torch.zeros(max_len, d_model) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe) def forward(self, x): return x self.pe[:x.size(1)]接下来组装微型Transformer编码器class MiniBERT(nn.Module): def __init__(self): super().__init__() self.embedding nn.Embedding(VOCAB_SIZE, EMBED_DIM) self.pos_encoder PositionalEncoding(EMBED_DIM) encoder_layer nn.TransformerEncoderLayer( d_modelEMBED_DIM, nheadN_HEADS) self.transformer nn.TransformerEncoder(encoder_layer, N_LAYERS) self.fc nn.Linear(EMBED_DIM, VOCAB_SIZE) def forward(self, src, maskNone): src self.embedding(src) * math.sqrt(EMBED_DIM) src self.pos_encoder(src) output self.transformer(src, mask) return self.fc(output)3. 实现Masked Language Model训练BERT的核心创新在于MLM预训练任务。我们实现一个简化的数据准备流程def create_masked_samples(text_tokens): 生成训练样本随机遮盖15%的token mask_prob 0.15 mask_token VOCAB_SIZE - 1 # 假设最后一个token是[MASK] masked_tokens text_tokens.clone() labels torch.full_like(text_tokens, -100) # 只计算被遮盖位置的loss # 随机选择要遮盖的位置 mask_positions torch.rand(text_tokens.shape) mask_prob # 80%替换为[MASK], 10%随机替换, 10%保持不变 labels[mask_positions] text_tokens[mask_positions] random_replace torch.rand(mask_positions.sum()) 0.1 random_tokens torch.randint(0, VOCAB_SIZE-1, (random_replace.sum(),)) masked_tokens[mask_positions] mask_token masked_tokens[mask_positions][random_replace] random_tokens return masked_tokens, labels训练循环的关键部分model MiniBERT() optimizer torch.optim.Adam(model.parameters(), lr1e-4) criterion nn.CrossEntropyLoss() for epoch in range(10): for batch in dataloader: inputs, labels create_masked_samples(batch) outputs model(inputs) loss criterion(outputs.view(-1, VOCAB_SIZE), labels.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step()4. 可视化注意力机制理解双向编码要真正理解BERT的双向性最好的方法是观察其注意力权重。我们提取并可视化第一个注意力头的权重import matplotlib.pyplot as plt def plot_attention(model, sentence): model.eval() tokens tokenize(sentence) src torch.LongTensor(tokens).unsqueeze(0) # 获取第一个Transformer层的注意力权重 with torch.no_grad(): output model.transformer.layers[0].self_attn( model.pos_encoder(model.embedding(src) * math.sqrt(EMBED_DIM)), model.pos_encoder(model.embedding(src) * math.sqrt(EMBED_DIM)), model.pos_encoder(model.embedding(src) * math.sqrt(EMBED_DIM)) )[1] # 返回注意力权重 plt.imshow(output.squeeze().numpy(), cmaphot) plt.xticks(range(len(tokens)), tokens) plt.yticks(range(len(tokens)), tokens) plt.show()运行plot_attention(model, the cat sat on the mat)您将看到每个单词如何关注句子中的其他单词这正是双向编码的直观体现。5. 进阶技巧从简化版到生产级BERT虽然我们的微型BERT只有不到50行代码但已经包含了BERT的核心思想。要将其发展为实用模型还需要更大规模的训练数据使用Wikipedia、BookCorpus等真实语料完整的分词系统实现WordPiece或SentencePiece分词更深的网络结构增加Transformer层数和注意力头数多任务学习加入Next Sentence Prediction任务优化技巧使用混合精度训练、梯度累积等# 生产级BERT的典型配置 class BERTConfig: vocab_size 30522 # WordPiece词表大小 hidden_size 768 # 隐藏层维度 num_hidden_layers 12 # Transformer层数 num_attention_heads 12 # 注意力头数 intermediate_size 3072 # FFN层维度 max_position_embeddings 512 # 最大位置编码6. 实际应用将微型BERT用于下游任务即使是我们的小模型也可以演示BERT的迁移学习能力。假设我们要做情感分析class SentimentClassifier(nn.Module): def __init__(self, bert_model): super().__init__() self.bert bert_model self.classifier nn.Linear(EMBED_DIM, 2) # 二分类 def forward(self, src): # 使用[CLS]位置的表示进行分类 output self.bert(src) cls_output output[:, 0, :] # 第一个位置是[CLS] return self.classifier(cls_output)微调时我们可以选择冻结BERT参数或联合训练# 加载预训练的微型BERT pretrained_model MiniBERT() pretrained_model.load_state_dict(torch.load(minibert.pth)) # 创建分类器并微调 classifier SentimentClassifier(pretrained_model) # 只训练分类头 for param in classifier.bert.parameters(): param.requires_grad False # 或者联合微调所有参数 # (需要更多数据和计算资源)在实现过程中我发现几个关键点对模型性能影响最大遮盖策略的随机性比例80-10-10规则位置编码的实现方式学习率设置和warmup策略注意力头数的选择要与嵌入维度匹配