用PyTorch从零手搓一个BERT模型从数据预处理到训练测试的保姆级教程在自然语言处理领域BERT无疑是一个里程碑式的模型。虽然现在我们可以轻松地通过Hugging Face等库调用预训练好的BERT但真正理解其内部机制的最佳方式莫过于亲手实现它。本文将带你从零开始用PyTorch一步步构建一个完整的BERT模型包括数据处理、模型架构实现、训练策略等核心环节。1. 环境准备与数据预处理在开始构建BERT之前我们需要准备好开发环境。建议使用Python 3.8和PyTorch 1.10版本。可以通过以下命令安装必要的依赖pip install torch numpy tqdmBERT的预训练需要大量文本数据。我们可以使用Wikipedia数据集或其他开源文本语料。数据预处理是BERT实现中最容易被忽视但至关重要的环节主要包括以下几个步骤文本清洗去除特殊字符、HTML标签等非文本内容分词处理使用WordPiece分词器将文本转换为子词单元特殊标记添加添加[CLS]、[SEP]等BERT特有的标记from transformers import BertTokenizer # 初始化分词器 tokenizer BertTokenizer.from_pretrained(bert-base-uncased) # 示例分词 text Implementing BERT from scratch is challenging but rewarding. tokens tokenizer.tokenize(text) print(tokens) # [implementing, bert, from, scratch, is, challenging, but, rewarding, .]2. BERT模型架构实现BERT的核心是Transformer编码器堆叠。让我们从最基本的组件开始构建。2.1 嵌入层实现BERT使用三种嵌入的组合词嵌入、位置嵌入和段嵌入。以下是嵌入层的PyTorch实现import torch import torch.nn as nn import math class BertEmbeddings(nn.Module): def __init__(self, vocab_size, hidden_size, max_position_embeddings, type_vocab_size, dropout_prob): super().__init__() self.word_embeddings nn.Embedding(vocab_size, hidden_size) self.position_embeddings nn.Embedding(max_position_embeddings, hidden_size) self.token_type_embeddings nn.Embedding(type_vocab_size, hidden_size) self.LayerNorm nn.LayerNorm(hidden_size) self.dropout nn.Dropout(dropout_prob) self.register_buffer(position_ids, torch.arange(max_position_embeddings).expand((1, -1))) def forward(self, input_ids, token_type_idsNone): 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) word_embeddings self.word_embeddings(input_ids) position_embeddings self.position_embeddings(position_ids) token_type_embeddings self.token_type_embeddings(token_type_ids) embeddings word_embeddings position_embeddings token_type_embeddings embeddings self.LayerNorm(embeddings) embeddings self.dropout(embeddings) return embeddings2.2 Transformer编码器层BERT的核心是多个Transformer编码器层的堆叠。每个编码器层包含自注意力机制和前馈网络class BertLayer(nn.Module): def __init__(self, hidden_size, num_attention_heads, intermediate_size, dropout_prob): super().__init__() self.attention BertAttention(hidden_size, num_attention_heads, dropout_prob) self.intermediate BertIntermediate(hidden_size, intermediate_size) self.output BertOutput(hidden_size, intermediate_size, dropout_prob) def forward(self, hidden_states, attention_maskNone): attention_output self.attention(hidden_states, attention_mask) intermediate_output self.intermediate(attention_output) layer_output self.output(intermediate_output, attention_output) return layer_output3. 预训练任务实现BERT通过两个预训练任务学习语言表示掩码语言模型(MLM)和下一句预测(NSP)。3.1 掩码语言模型(MLM)MLM随机掩盖输入token的15%并让模型预测被掩盖的token。实现这一任务有几个关键点随机选择15%的token进行掩盖其中80%替换为[MASK]10%替换为随机token10%保持不变使用交叉熵损失计算预测结果def create_masked_lm_predictions(tokens, mask_token_id, vocab_size, mask_prob0.15): 创建MLM任务的输入和标签 output_tokens list(tokens) masked_lm_positions [] masked_lm_labels [] for i, token in enumerate(tokens): if token in [cls_token_id, sep_token_id]: continue prob random.random() if prob mask_prob: masked_lm_positions.append(i) masked_lm_labels.append(token) # 80%概率替换为[MASK] if prob mask_prob * 0.8: output_tokens[i] mask_token_id # 10%概率替换为随机token elif mask_prob * 0.8 prob mask_prob * 0.9: output_tokens[i] random.randrange(vocab_size) # 10%概率保持不变 else: pass return output_tokens, masked_lm_positions, masked_lm_labels3.2 下一句预测(NSP)NSP任务判断两个句子是否是连续的。实现时需要注意50%的样本使用实际连续的句子对50%的样本使用随机组合的句子对使用二元交叉熵损失计算预测结果def create_next_sentence_predictions(sentence_a, sentence_b): 创建NSP任务的输入和标签 # 50%概率使用实际下一句 if random.random() 0.5: is_next 1 actual_next get_actual_next_sentence(sentence_a) tokens_b tokenizer.tokenize(actual_next) # 50%概率使用随机句子 else: is_next 0 random_sentence get_random_sentence() tokens_b tokenizer.tokenize(random_sentence) return tokens_b, is_next4. 训练策略与技巧训练BERT模型需要特别注意以下几个关键点4.1 学习率调度BERT使用带warmup的学习率调度策略这对于训练稳定性至关重要def get_learning_rate(step, warmup_steps, initial_lr): 带warmup的学习率调度 if step warmup_steps: return initial_lr * (step / warmup_steps) else: return initial_lr * (warmup_steps ** 0.5) * (step ** -0.5)4.2 梯度累积由于BERT模型参数量大通常需要使用梯度累积技术optimizer.zero_grad() for i, batch in enumerate(train_dataloader): loss model(**batch) loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()4.3 混合精度训练使用混合精度训练可以显著减少显存占用并加速训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss model(**batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5. 模型评估与调试训练过程中需要监控多个指标来评估模型性能指标名称计算方法预期范围MLM准确率预测正确的masked token逐渐提高NSP准确率预测正确的句子关系50%训练损失平均batch损失逐渐降低验证损失验证集上的平均损失低于训练集调试BERT模型时常见的几个问题损失不下降检查学习率是否合适数据预处理是否正确梯度爆炸尝试梯度裁剪调整学习率过拟合增加dropout率使用更多训练数据# 示例评估代码 def evaluate(model, eval_dataloader): model.eval() total_loss 0 total_correct 0 total_samples 0 with torch.no_grad(): for batch in eval_dataloader: outputs model(**batch) total_loss outputs.loss.item() total_correct (outputs.logits.argmax(-1) batch[labels]).sum().item() total_samples batch[labels].numel() avg_loss total_loss / len(eval_dataloader) accuracy total_correct / total_samples return avg_loss, accuracy6. 实际应用与优化完成预训练后我们可以将BERT模型应用于下游任务。以下是一些优化技巧动态掩码每次epoch重新生成掩码模式增加数据多样性全词掩码对完整单词进行掩码而非单独的子词梯度检查点减少显存占用允许更大的batch size# 动态掩码实现示例 class DynamicMaskingDataset(torch.utils.data.Dataset): def __getitem__(self, index): text self.texts[index] tokens self.tokenizer.tokenize(text) # 每次获取数据时重新生成掩码 masked_tokens, masked_positions, masked_labels create_masked_lm_predictions(tokens) return { input_ids: torch.tensor(self.tokenizer.convert_tokens_to_ids(masked_tokens)), masked_lm_positions: torch.tensor(masked_positions), masked_lm_labels: torch.tensor(masked_labels) }实现BERT模型从零开始确实是一项挑战但通过这个过程你将对Transformer架构和预训练语言模型有更深入的理解。在实际项目中我发现在模型规模较小时适当增加训练步数比单纯扩大模型规模更有效。另外合理的数据预处理和清洗往往比模型结构调整带来的提升更明显。
用PyTorch从零手搓一个BERT模型:从数据预处理到训练测试的保姆级教程
用PyTorch从零手搓一个BERT模型从数据预处理到训练测试的保姆级教程在自然语言处理领域BERT无疑是一个里程碑式的模型。虽然现在我们可以轻松地通过Hugging Face等库调用预训练好的BERT但真正理解其内部机制的最佳方式莫过于亲手实现它。本文将带你从零开始用PyTorch一步步构建一个完整的BERT模型包括数据处理、模型架构实现、训练策略等核心环节。1. 环境准备与数据预处理在开始构建BERT之前我们需要准备好开发环境。建议使用Python 3.8和PyTorch 1.10版本。可以通过以下命令安装必要的依赖pip install torch numpy tqdmBERT的预训练需要大量文本数据。我们可以使用Wikipedia数据集或其他开源文本语料。数据预处理是BERT实现中最容易被忽视但至关重要的环节主要包括以下几个步骤文本清洗去除特殊字符、HTML标签等非文本内容分词处理使用WordPiece分词器将文本转换为子词单元特殊标记添加添加[CLS]、[SEP]等BERT特有的标记from transformers import BertTokenizer # 初始化分词器 tokenizer BertTokenizer.from_pretrained(bert-base-uncased) # 示例分词 text Implementing BERT from scratch is challenging but rewarding. tokens tokenizer.tokenize(text) print(tokens) # [implementing, bert, from, scratch, is, challenging, but, rewarding, .]2. BERT模型架构实现BERT的核心是Transformer编码器堆叠。让我们从最基本的组件开始构建。2.1 嵌入层实现BERT使用三种嵌入的组合词嵌入、位置嵌入和段嵌入。以下是嵌入层的PyTorch实现import torch import torch.nn as nn import math class BertEmbeddings(nn.Module): def __init__(self, vocab_size, hidden_size, max_position_embeddings, type_vocab_size, dropout_prob): super().__init__() self.word_embeddings nn.Embedding(vocab_size, hidden_size) self.position_embeddings nn.Embedding(max_position_embeddings, hidden_size) self.token_type_embeddings nn.Embedding(type_vocab_size, hidden_size) self.LayerNorm nn.LayerNorm(hidden_size) self.dropout nn.Dropout(dropout_prob) self.register_buffer(position_ids, torch.arange(max_position_embeddings).expand((1, -1))) def forward(self, input_ids, token_type_idsNone): 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) word_embeddings self.word_embeddings(input_ids) position_embeddings self.position_embeddings(position_ids) token_type_embeddings self.token_type_embeddings(token_type_ids) embeddings word_embeddings position_embeddings token_type_embeddings embeddings self.LayerNorm(embeddings) embeddings self.dropout(embeddings) return embeddings2.2 Transformer编码器层BERT的核心是多个Transformer编码器层的堆叠。每个编码器层包含自注意力机制和前馈网络class BertLayer(nn.Module): def __init__(self, hidden_size, num_attention_heads, intermediate_size, dropout_prob): super().__init__() self.attention BertAttention(hidden_size, num_attention_heads, dropout_prob) self.intermediate BertIntermediate(hidden_size, intermediate_size) self.output BertOutput(hidden_size, intermediate_size, dropout_prob) def forward(self, hidden_states, attention_maskNone): attention_output self.attention(hidden_states, attention_mask) intermediate_output self.intermediate(attention_output) layer_output self.output(intermediate_output, attention_output) return layer_output3. 预训练任务实现BERT通过两个预训练任务学习语言表示掩码语言模型(MLM)和下一句预测(NSP)。3.1 掩码语言模型(MLM)MLM随机掩盖输入token的15%并让模型预测被掩盖的token。实现这一任务有几个关键点随机选择15%的token进行掩盖其中80%替换为[MASK]10%替换为随机token10%保持不变使用交叉熵损失计算预测结果def create_masked_lm_predictions(tokens, mask_token_id, vocab_size, mask_prob0.15): 创建MLM任务的输入和标签 output_tokens list(tokens) masked_lm_positions [] masked_lm_labels [] for i, token in enumerate(tokens): if token in [cls_token_id, sep_token_id]: continue prob random.random() if prob mask_prob: masked_lm_positions.append(i) masked_lm_labels.append(token) # 80%概率替换为[MASK] if prob mask_prob * 0.8: output_tokens[i] mask_token_id # 10%概率替换为随机token elif mask_prob * 0.8 prob mask_prob * 0.9: output_tokens[i] random.randrange(vocab_size) # 10%概率保持不变 else: pass return output_tokens, masked_lm_positions, masked_lm_labels3.2 下一句预测(NSP)NSP任务判断两个句子是否是连续的。实现时需要注意50%的样本使用实际连续的句子对50%的样本使用随机组合的句子对使用二元交叉熵损失计算预测结果def create_next_sentence_predictions(sentence_a, sentence_b): 创建NSP任务的输入和标签 # 50%概率使用实际下一句 if random.random() 0.5: is_next 1 actual_next get_actual_next_sentence(sentence_a) tokens_b tokenizer.tokenize(actual_next) # 50%概率使用随机句子 else: is_next 0 random_sentence get_random_sentence() tokens_b tokenizer.tokenize(random_sentence) return tokens_b, is_next4. 训练策略与技巧训练BERT模型需要特别注意以下几个关键点4.1 学习率调度BERT使用带warmup的学习率调度策略这对于训练稳定性至关重要def get_learning_rate(step, warmup_steps, initial_lr): 带warmup的学习率调度 if step warmup_steps: return initial_lr * (step / warmup_steps) else: return initial_lr * (warmup_steps ** 0.5) * (step ** -0.5)4.2 梯度累积由于BERT模型参数量大通常需要使用梯度累积技术optimizer.zero_grad() for i, batch in enumerate(train_dataloader): loss model(**batch) loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()4.3 混合精度训练使用混合精度训练可以显著减少显存占用并加速训练scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): loss model(**batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5. 模型评估与调试训练过程中需要监控多个指标来评估模型性能指标名称计算方法预期范围MLM准确率预测正确的masked token逐渐提高NSP准确率预测正确的句子关系50%训练损失平均batch损失逐渐降低验证损失验证集上的平均损失低于训练集调试BERT模型时常见的几个问题损失不下降检查学习率是否合适数据预处理是否正确梯度爆炸尝试梯度裁剪调整学习率过拟合增加dropout率使用更多训练数据# 示例评估代码 def evaluate(model, eval_dataloader): model.eval() total_loss 0 total_correct 0 total_samples 0 with torch.no_grad(): for batch in eval_dataloader: outputs model(**batch) total_loss outputs.loss.item() total_correct (outputs.logits.argmax(-1) batch[labels]).sum().item() total_samples batch[labels].numel() avg_loss total_loss / len(eval_dataloader) accuracy total_correct / total_samples return avg_loss, accuracy6. 实际应用与优化完成预训练后我们可以将BERT模型应用于下游任务。以下是一些优化技巧动态掩码每次epoch重新生成掩码模式增加数据多样性全词掩码对完整单词进行掩码而非单独的子词梯度检查点减少显存占用允许更大的batch size# 动态掩码实现示例 class DynamicMaskingDataset(torch.utils.data.Dataset): def __getitem__(self, index): text self.texts[index] tokens self.tokenizer.tokenize(text) # 每次获取数据时重新生成掩码 masked_tokens, masked_positions, masked_labels create_masked_lm_predictions(tokens) return { input_ids: torch.tensor(self.tokenizer.convert_tokens_to_ids(masked_tokens)), masked_lm_positions: torch.tensor(masked_positions), masked_lm_labels: torch.tensor(masked_labels) }实现BERT模型从零开始确实是一项挑战但通过这个过程你将对Transformer架构和预训练语言模型有更深入的理解。在实际项目中我发现在模型规模较小时适当增加训练步数比单纯扩大模型规模更有效。另外合理的数据预处理和清洗往往比模型结构调整带来的提升更明显。