从零开始预训练BERT模型

import torch
import torch.nn as nn
import torch.nn.functional as F

# ... 其他导入

def create_mlm_mask(input_tensor, mask_ratio=0.15):
    """随机掩盖输入张量中的Token"""
    batch_size, seq_len = input_tensor.shape
    mask = torch.zeros_like(input_tensor, dtype=torch.bool)
    
    # 计算需要掩盖的Token数量
    num_to_mask = int(seq_len * mask_ratio)
    
    # 为每个样本生成掩码
    for i in range(batch_size):
        # 随机选择需要掩盖的Token索引（不包括特殊Token，如[CLS]和[SEP]）
        # 假设[CLS]在索引0，[SEP]在索引seq_len-1
        token_indices = torch.randperm(seq_len - 2)[:num_to_mask] + 1
        mask[i, token_indices] = True
        
    # 假设我们有一个特殊的MASK Token ID，比如103
    MASK_TOKEN_ID = 103
    masked_input = input_tensor.clone()
    masked_input[mask] = MASK_TOKEN_ID
    
    return masked_input, mask

class TinyBERT(nn.Module):
    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=2, max_len=512, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        
        # 嵌入层
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)

        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=d_model*4, dropout=dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # MLM头：将模型输出映射回词汇表大小
        self.mlm_head = nn.Linear(d_model, vocab_size)

    def forward(self, input_ids):
        # 1. 创建嵌入
        token_ids = input_ids
        batch_size, seq_len = token_ids.shape
        
        # 创建位置索引 [0, 1, ..., seq_len-1]
        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0).repeat(batch_size, 1)
        
        # 计算嵌入
        token_emb = self.token_embedding(token_ids)
        pos_emb = self.position_embedding(positions)
        
        # 组合嵌入
        embeddings = token_emb + pos_emb
        embeddings = self.dropout(embeddings)
        
        # 2. Transformer编码器
        # 注意：这里我们假设batch_first=True，如果使用旧版PyTorch可能需要转置
        transformer_output = self.transformer_encoder(embeddings)
        
        # 3. MLM预测
        mlm_logits = self.mlm_head(transformer_output)
        
        return mlm_logits

# --- 假设的参数和模型初始化 ---
VOCAB_SIZE = 30000
MAX_LEN = 128

model = TinyBERT(vocab_size=VOCAB_SIZE, d_model=128, num_layers=2, max_len=MAX_LEN)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss(ignore_index=0) # 假设0是填充（padding）Token ID

# --- 模拟数据加载器 (DataLoader) ---
# 假设我们有一个包含Token ID序列的张量
# 真实场景中，数据加载器会处理数据批次和填充

def get_dummy_batch(batch_size=4, seq_len=MAX_LEN):
    # 模拟一些随机的输入Token ID（避免使用特殊Token ID 0, 103）
    return torch.randint(low=2, high=VOCAB_SIZE, size=(batch_size, seq_len))

# --- 训练步骤 ---
model.train()

input_ids = get_dummy_batch()

# 1. 创建MLM掩码
masked_input, mlm_mask = create_mlm_mask(input_ids, mask_ratio=0.15)

# 2. 前向传播
mlm_logits = model(masked_input)

# 3. 计算损失
# 仅计算被掩盖Token的损失
# 维度: [B*L, V] vs [B*L]

# 将logits和掩码展平，只保留被掩盖的Token位置
logits_for_loss = mlm_logits[mlm_mask]
labels_for_loss = input_ids[mlm_mask]

# 检查是否有任何Token被掩盖
if labels_for_loss.numel() > 0:
    loss = criterion(logits_for_loss, labels_for_loss)
    
    # 4. 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    print(f"MLM Loss: {loss.item():.4f}")
else:
    print("Warning: No tokens were masked in this batch.")

# ------------------------------------------------------------------
# 模拟NSP任务（此处仅为概念性展示，实际实现需要更复杂的输入和输出处理）

class BertForNSP(nn.Module):
    def __init__(self, bert_model):
        super().__init__()
        self.bert = bert_model
        # NSP头使用[CLS]Token (索引0)的输出进行预测
        self.nsp_head = nn.Linear(bert_model.d_model, 2) # 2 classes: IsNext or NotNext
        
    def forward(self, input_ids, segment_ids=None):
        # 假设我们已经处理了输入，并获取了所有层的输出
        transformer_output = self.bert.transformer_encoder(self.bert.embeddings(input_ids)) # 简化处理
        
        # 使用[CLS] Token的输出 (通常是第一个Token的输出)
        cls_output = transformer_output[:, 0, :]
        nsp_logits = self.nsp_head(cls_output)
        return nsp_logits

# ------------------------------------------------------------------