别再死记硬背BERT了！用Python代码手撕Masked LM和Next Sentence Prediction

从零实现BERT核心预训练任务Masked LM与Next Sentence Prediction实战指南如果你已经厌倦了反复背诵BERT的八股文理论不妨跟随本文一起动手用Python代码揭开BERT预训练任务的神秘面纱。本文将带你从零开始实现Masked Language ModelMLM和Next Sentence PredictionNSP两大核心任务通过代码实践深入理解BERT的工作原理。1. 环境准备与数据加载在开始之前我们需要搭建一个适合深度学习实验的环境。推荐使用Python 3.8和PyTorch 1.8版本这些版本在稳定性和性能方面都有不错的表现。首先安装必要的依赖库pip install torch transformers numpy tqdm对于本实验我们将使用一个小型文本数据集来演示BERT的预训练过程。在实际应用中你可以替换为更大规模的语料库。import torch from torch.utils.data import Dataset class TextDataset(Dataset): def __init__(self, texts, max_length128): self.texts texts self.max_length max_length def __len__(self): return len(self.texts) def __getitem__(self, idx): text self.texts[idx] # 这里简化处理实际应用中应添加分词等预处理 tokens text.split()[:self.max_length] return tokens # 示例数据 sample_texts [ 深度学习是机器学习的一个分支, BERT是一种基于Transformer的预训练模型, 自然语言处理是人工智能的重要领域, 预训练语言模型在多项NLP任务中表现出色 ] dataset TextDataset(sample_texts)2. 实现Masked Language ModelMasked Language Model是BERT的核心预训练任务之一它通过随机遮盖输入文本中的部分token然后让模型预测这些被遮盖的原始token。2.1 理解MLM的遮盖策略BERT采用了一种特殊的遮盖策略随机选择15%的token进行处理其中80%被替换为[MASK]标记10%被替换为随机token10%保持不变这种策略的设计有两个主要目的使模型具备一定的文本纠错能力缓解预训练与微调阶段的输入不匹配问题2.2 实现遮盖逻辑下面我们实现一个MLM数据处理器import random from collections import defaultdict class MLMProcessor: def __init__(self, vocab, mask_token[MASK], mask_prob0.15): self.vocab vocab self.mask_token mask_token self.mask_prob mask_prob self.vocab_size len(vocab) def mask_tokens(self, tokens): masked_tokens tokens.copy() labels [-100] * len(tokens) # -100表示不计算损失 for i, token in enumerate(tokens): if random.random() self.mask_prob: labels[i] self.vocab.get(token, self.vocab[[UNK]]) rand random.random() if rand 0.8: masked_tokens[i] self.mask_token elif rand 0.9: masked_tokens[i] random.choice(list(self.vocab.keys())) # 剩下10%保持不变 return masked_tokens, labels # 示例词汇表 vocab { [PAD]: 0, [UNK]: 1, [CLS]: 2, [SEP]: 3, [MASK]: 4, 深度: 5, 学习: 6, 是: 7, 机器: 8, 一个: 9, 分支: 10, BERT: 11, 一种: 12, 基于: 13, Transformer: 14, 的: 15, 预训练: 16, 模型: 17, 自然: 18, 语言: 19, 处理: 20, 人工: 21, 智能: 22, 重要: 23, 领域: 24, 在: 25, 多项: 26, NLP: 27, 任务: 28, 中: 29, 表现: 30, 出色: 31 } mlm_processor MLMProcessor(vocab) tokens BERT 是一种基于 Transformer 的预训练模型.split() masked, labels mlm_processor.mask_tokens(tokens) print(原始tokens:, tokens) print(遮盖后:, masked) print(标签:, labels)2.3 构建MLM模型现在我们来构建一个简化的BERT模型专注于MLM任务import torch.nn as nn import torch.nn.functional as F class SimpleBERT(nn.Module): def __init__(self, vocab_size, hidden_size768, num_layers2, num_heads2): super().__init__() self.embedding nn.Embedding(vocab_size, hidden_size) self.encoder_layers nn.ModuleList([ nn.TransformerEncoderLayer(hidden_size, num_heads, hidden_size*4) for _ in range(num_layers) ]) self.lm_head nn.Linear(hidden_size, vocab_size) def forward(self, input_ids): x self.embedding(input_ids) for layer in self.encoder_layers: x layer(x) logits self.lm_head(x) return logits model SimpleBERT(len(vocab)) print(model)3. 实现Next Sentence PredictionNext Sentence Prediction是BERT的另一个重要预训练任务它让模型判断两个句子是否是连续的。3.1 理解NSP任务NSP的设计目的是让模型理解句子间的关系这对问答、自然语言推理等任务非常重要。在预训练中50%的情况下第二个句子是第一个句子的实际后续50%的情况下第二个句子是从语料库中随机选取的3.2 构建NSP数据集我们需要构建包含句子对的数据并标注它们是否是连续的class NSPDataset(Dataset): def __init__(self, sentence_pairs): self.pairs sentence_pairs def __len__(self): return len(self.pairs) def __getitem__(self, idx): pair self.pairs[idx] return pair[sent1], pair[sent2], pair[is_next] # 示例数据 nsp_pairs [ { sent1: 深度学习是机器学习的一个分支, sent2: 它通过多层神经网络学习数据表示, is_next: 1 # 是连续句子 }, { sent1: BERT是一种预训练模型, sent2: 太阳系有八大行星, is_next: 0 # 不是连续句子 } ] nsp_dataset NSPDataset(nsp_pairs)3.3 扩展BERT模型支持NSP我们需要扩展之前的SimpleBERT模型使其能够处理NSP任务class BERTWithNSP(SimpleBERT): def __init__(self, vocab_size, hidden_size768): super().__init__(vocab_size, hidden_size) self.cls_token_id vocab[[CLS]] self.sep_token_id vocab[[SEP]] self.nsp_head nn.Linear(hidden_size, 2) def forward(self, input_ids, segment_idsNone): # 添加[CLS]和[SEP]标记 batch_size input_ids.size(0) cls_tokens torch.full((batch_size, 1), self.cls_token_id, deviceinput_ids.device) sep_tokens torch.full((batch_size, 1), self.sep_token_id, deviceinput_ids.device) # 构建输入 (简化版实际BERT输入更复杂) input_embeddings self.embedding(input_ids) # 通过Transformer编码器 x input_embeddings for layer in self.encoder_layers: x layer(x) # [CLS]位置的输出用于NSP cls_output x[:, 0, :] nsp_logits self.nsp_head(cls_output) # 其余位置输出用于MLM lm_logits self.lm_head(x) return lm_logits, nsp_logits bert_model BERTWithNSP(len(vocab)) print(bert_model)4. 联合训练与可视化现在我们将MLM和NSP任务结合起来进行联合训练并可视化训练过程。4.1 训练循环实现from tqdm import tqdm import matplotlib.pyplot as plt def train(model, mlm_processor, dataset, epochs10, lr1e-4): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) optimizer torch.optim.Adam(model.parameters(), lrlr) mlm_losses [] nsp_losses [] for epoch in range(epochs): model.train() total_mlm_loss 0 total_nsp_loss 0 for text in tqdm(dataset, descfEpoch {epoch1}): # 简化处理实际应用中应对文本进行更复杂的预处理 tokens text masked_tokens, mlm_labels mlm_processor.mask_tokens(tokens) # 转换为tensor input_ids torch.tensor([vocab.get(t, vocab[[UNK]]) for t in masked_tokens]).unsqueeze(0).to(device) mlm_labels torch.tensor(mlm_labels).unsqueeze(0).to(device) # 随机生成NSP标签 (简化版) is_next torch.tensor([random.randint(0, 1)]).to(device) # 前向传播 lm_logits, nsp_logits model(input_ids) # 计算损失 mlm_loss F.cross_entropy( lm_logits.view(-1, len(vocab)), mlm_labels.view(-1), ignore_index-100 ) nsp_loss F.cross_entropy( nsp_logits.view(-1, 2), is_next.view(-1) ) total_loss mlm_loss nsp_loss # 反向传播 optimizer.zero_grad() total_loss.backward() optimizer.step() total_mlm_loss mlm_loss.item() total_nsp_loss nsp_loss.item() avg_mlm_loss total_mlm_loss / len(dataset) avg_nsp_loss total_nsp_loss / len(dataset) mlm_losses.append(avg_mlm_loss) nsp_losses.append(avg_nsp_loss) print(fEpoch {epoch1}: MLM Loss {avg_mlm_loss:.4f}, NSP Loss {avg_nsp_loss:.4f}) # 绘制损失曲线 plt.figure(figsize(10, 5)) plt.plot(mlm_losses, labelMLM Loss) plt.plot(nsp_losses, labelNSP Loss) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() plt.title(Training Loss) plt.show() # 开始训练 train(bert_model, mlm_processor, dataset, epochs20)4.2 结果分析与调试技巧训练完成后我们可以通过一些示例来验证模型的表现def predict_masked_token(model, text, mask_position): device next(model.parameters()).device tokens text.split() tokens[mask_position] [MASK] input_ids torch.tensor([vocab.get(t, vocab[[UNK]]) for t in tokens]).unsqueeze(0).to(device) with torch.no_grad(): lm_logits, _ model(input_ids) predictions lm_logits[0, mask_position] probs F.softmax(predictions, dim-1) top_k 3 top_probs, top_indices torch.topk(probs, top_k) print(f原始文本: {text}) print(f遮盖后: { .join(tokens)}) print(预测结果:) for i in range(top_k): token list(vocab.keys())[list(vocab.values()).index(top_indices[i].item())] print(f{token}: {top_probs[i].item():.4f}) # 测试MLM test_text BERT 是一种预训练模型 predict_masked_token(bert_model, test_text, mask_position3) def predict_next_sentence(model, sent1, sent2): device next(model.parameters()).device tokens [[CLS]] sent1.split() [[SEP]] sent2.split() [[SEP]] input_ids torch.tensor([vocab.get(t, vocab[[UNK]]) for t in tokens]).unsqueeze(0).to(device) with torch.no_grad(): _, nsp_logits model(input_ids) probs F.softmax(nsp_logits, dim-1) print(f句子1: {sent1}) print(f句子2: {sent2}) print(f是下一句的概率: {probs[0, 1].item():.4f}) print(f不是下一句的概率: {probs[0, 0].item():.4f}) # 测试NSP sent1 深度学习是机器学习的一个分支 sent2 它通过多层神经网络学习数据表示 predict_next_sentence(bert_model, sent1, sent2) sent1 BERT是一种预训练模型 sent2 太阳系有八大行星 predict_next_sentence(bert_model, sent1, sent2)在实际应用中你可能会遇到各种问题这里提供一些调试建议损失不下降检查学习率是否合适确认数据预处理是否正确尝试减小模型规模过拟合增加训练数据量添加Dropout层使用更小的模型训练不稳定尝试梯度裁剪使用学习率预热检查数据中的异常值5. 进阶优化与扩展现在我们已经实现了BERT的核心预训练任务接下来可以探索一些进阶优化和扩展方向。5.1 使用Hugging Face Transformers库虽然我们从零开始实现了BERT的核心功能但在实际项目中我们通常会使用Hugging Face的Transformers库from transformers import BertTokenizer, BertForPreTraining import torch tokenizer BertTokenizer.from_pretrained(bert-base-chinese) model BertForPreTraining.from_pretrained(bert-base-chinese) inputs tokenizer(深度学习是[MASK]学习的一个分支, return_tensorspt) outputs model(**inputs) predictions outputs.prediction_logits masked_index (inputs.input_ids tokenizer.mask_token_id)[0].nonzero().item() predicted_token tokenizer.convert_ids_to_tokens(predictions[0, masked_index].argmax().item()) print(f预测结果: {predicted_token})5.2 全词掩码(Whole Word Masking)实现原始的BERT使用token级别的掩码而全词掩码会掩码整个词而不仅仅是子词class WWMProcessor(MLMProcessor): def __init__(self, vocab, tokenizer, mask_prob0.15): super().__init__(vocab, mask_probmask_prob) self.tokenizer tokenizer def mask_whole_words(self, text): words text.split() masked_words words.copy() labels [-100] * len(words) # 先分词获取单词到token的映射 word_to_tokens [] for word in words: tokens self.tokenizer.tokenize(word) word_to_tokens.append((word, tokens)) # 实施全词掩码 for i, (word, tokens) in enumerate(word_to_tokens): if random.random() self.mask_prob: labels[i] self.vocab.get(word, self.vocab[[UNK]]) rand random.random() if rand 0.8: masked_words[i] self.mask_token elif rand 0.9: masked_words[i] random.choice(list(self.vocab.keys())) # 剩下10%保持不变 return .join(masked_words), labels # 示例使用 wwm_processor WWMProcessor(vocab, lambda x: x.split()) # 简化分词器 text BERT 是一种预训练模型 masked, labels wwm_processor.mask_whole_words(text) print(原始:, text) print(掩码后:, masked) print(标签:, labels)5.3 多GPU训练与混合精度对于大规模预训练我们需要利用多GPU和混合精度训练来加速from torch.nn.parallel import DataParallel from torch.cuda.amp import GradScaler, autocast def train_with_amp(model, dataset, epochs10): device torch.device(cuda if torch.cuda.is_available() else cpu) model DataParallel(model).to(device) optimizer torch.optim.Adam(model.parameters(), lr1e-4) scaler GradScaler() for epoch in range(epochs): model.train() total_loss 0 for text in tqdm(dataset, descfEpoch {epoch1}): optimizer.zero_grad() with autocast(): # 这里简化了输入处理 tokens text masked_tokens, mlm_labels mlm_processor.mask_tokens(tokens) input_ids torch.tensor([vocab.get(t, vocab[[UNK]]) for t in masked_tokens]).unsqueeze(0).to(device) mlm_labels torch.tensor(mlm_labels).unsqueeze(0).to(device) is_next torch.tensor([random.randint(0, 1)]).to(device) lm_logits, nsp_logits model(input_ids) mlm_loss F.cross_entropy( lm_logits.view(-1, len(vocab)), mlm_labels.view(-1), ignore_index-100 ) nsp_loss F.cross_entropy( nsp_logits.view(-1, 2), is_next.view(-1) ) loss mlm_loss nsp_loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss loss.item() print(fEpoch {epoch1}: Loss {total_loss/len(dataset):.4f}) # 使用混合精度训练 train_with_amp(bert_model, dataset)通过本文的实践我们不仅理解了BERT预训练的核心机制还掌握了如何从零开始实现这些功能。这种深入的理解将帮助你在实际应用中更好地微调和优化BERT模型。