BERT预训练实战：从零开始构建你的第一个掩码语言模型（附代码）-CSDN博客

从零构建你的BERT：一次彻底搞懂预训练与微调的实战指南

如果你刚踏入自然语言处理（NLP）领域，面对BERT、RoBERTa、DeBERTa这些层出不穷的模型感到眼花缭乱，或者你虽然调过Hugging Face的from_pretrained，但对模型内部究竟发生了什么仍心存疑惑，那么这篇文章正是为你准备的。我们不会停留在理论层面的高谈阔论，而是直接动手，从最原始的文本数据开始，一步步构建一个能够进行掩码语言建模（MLM）的微型BERT，并探讨如何将其适配到你的具体任务上。这不仅仅是一次代码演练，更是一次深入理解Transformer架构及其训练范式的思维之旅。对于中高级开发者而言，文中涉及的工程技巧和设计权衡，或许也能带来新的启发。

我们将使用PyTorch和Hugging Face的transformers库作为主要工具，但核心逻辑会剥离框架，让你看清本质。整个流程将涵盖数据准备、Tokenizer原理、模型架构实现、预训练任务设计、训练循环优化以及下游微调策略。准备好了吗？让我们开始这场从零到一的模型创造之旅。

1. 环境准备与数据工程

在开始构建模型之前，一个稳定且高效的开发环境至关重要。我们选择Google Colab作为实验平台，它提供了免费的GPU资源，非常适合进行中等规模的语言模型预训练实验。

1.1 核心依赖安装

首先，我们需要安装必要的Python库。除了标准的科学计算栈，transformers和datasets库是我们的左膀右臂。

!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
!pip install transformers datasets accelerate sentencepiece
!pip install wandb  # 用于实验跟踪，可选但强烈推荐

安装完成后，导入我们将要使用的核心模块：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, DataCollatorForLanguageModeling
from datasets import load_dataset
import numpy as np
import random
from tqdm.auto import tqdm

提示：使用accelerate库可以轻松实现混合精度训练和分布式训练，对于更大规模的预训练能显著提升速度并节省显存。

1.2 构建一个高效的文本数据集

预训练需要海量的无标注文本。这里我们使用datasets库加载一个经典的中文语料库——维基百科中文版。关键在于，我们需要将原始文本处理成模型能够消化的小块。

# 加载数据集
raw_dataset = load_dataset("pleisto/wikipedia-cn-20230720-filtered", split="train[:1%]") # 仅使用1%作为演示
print(f"数据集大小: {len(raw_dataset)} 条文本")

# 定义一个简单的文本处理函数
def process_function(examples):
    # 移除过短的行和空白字符
    texts = [text.strip() for text in examples['text'] if len(text.strip()) > 10]
    return {'text': texts}

processed_dataset = raw_dataset.map(process_function, batched=True, remove_columns=raw_dataset.column_names)

原始文本是长篇文章，而BERT的输入长度有限（通常是512个token）。我们需要进行“文档分割”（Document Sharding）。一种实用的策略是使用滑动窗口，将长文本切分成固定长度的片段，并允许一定的重叠以避免在片段边界处切断重要信息。

class TextChunkDataset(Dataset):
    def __init__(self, text_list, tokenizer, max_length=512, stride=128):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.stride = stride
        self.chunks = []

        for text in tqdm(text_list, desc="处理文本块"):
            # 将文本编码为token IDs
            tokens = tokenizer.encode(text, add_special_tokens=False)
            # 使用滑动窗口创建块
            for i in range(0, len(tokens), stride):
                chunk = tokens[i: i + max_length]
                if len(chunk) < 16:  # 忽略过短的块
                    continue
                # 填充或截断到max_length
                if len(chunk) < max_length:
                    chunk = chunk + [tokenizer.pad_token_id] * (max_length - len(chunk))
                else:
                    chunk = chunk[:max_length]
                self.chunks.append(chunk)

    def __len__(self):
        return len(self.chunks)

    def __getitem__(self, idx):
        return torch.tensor(self.chunks[idx], dtype=torch.long)

这个TextChunkDataset类完成了从原始文本到模型可输入张量的转换。stride参数控制了滑动窗口的步长，较小的步长会产生更多重叠的片段，增加了数据的多样性，但也会增大数据集规模。

2. 解剖BERT：亲手实现核心组件

理解了数据流，接下来我们深入到模型内部。我们将实现一个简化版的BERT，专注于其最核心的部分：嵌入层、Transformer编码器和MLM头。这能让你彻底摆脱“黑箱”感。

2.1 嵌入层：从符号到向量的桥梁

BERT的嵌入层由三个部分组成：词元嵌入（Token Embeddings）、段嵌入（Segment Embeddings）和位置嵌入（Position Embeddings）。在预训练中，段嵌入用于下一句预测（NSP）任务，但在我们专注于MLM的简化版中，可以暂时忽略。

class BERTEmbeddings(nn.Module):
    def __init__(self, vocab_size, hidden_size=768, max_position_embeddings=512, layer_norm_eps=1e-12, hidden_dropout_prob=0.1):
        super().__init__()
        self.token_embeddings = nn.Embedding(vocab_size, hidden_size)
        self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size)
        # 在仅MLM任务中，我们简化处理，不使用段嵌入
        # self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size)

        self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
        self.dropout = nn.Dropout(hidd