4、数据组件（二）：数据加载器、批次生成、填充与掩码

好，咱们接着聊数据组件。上一章我们把数据集和词表搞定了，数据已经变成了整齐的 token ID 序列。但有个问题——这些序列长度不一样。

你想想看，有的句子 5 个词，有的 20 个词。Transformer 内部是矩阵运算，要求所有输入形状一致。怎么办？

说白了就两个动作：填充和掩码。

4.1 批次生成——把样本打包

我个人习惯把批次生成放在数据加载器里一起做。PyTorch 的 DataLoader 有个参数叫 collate_fn，它就是干这个活的。

先看一个最简单的批次生成逻辑：

def collate_batch(batch):
    # batch 是一个列表，每个元素是 (input_ids, target_ids)
    inputs = [item[0] for item in batch]
    targets = [item[1] for item in batch]
    
    # 找到当前批次里最长的序列
    max_len = max(len(seq) for seq in inputs)
    
    # 填充到相同长度
    padded_inputs = []
    padded_targets = []
    
    for inp, tgt in zip(inputs, targets):
        pad_len = max_len - len(inp)
        # 用 0 填充（假设 0 是 pad_token_id）
        padded_inp = inp + [0] * pad_len
        padded_tgt = tgt + [0] * pad_len
        padded_inputs.append(padded_inp)
        padded_targets.append(padded_tgt)
    
    return torch.tensor(padded_inputs), torch.tensor(padded_targets)

嗯，这里要注意：填充值必须和词表中的 <pad> token 对应。我在项目中遇到过有人用 0 填充，但词表的 <pad> 是 1，结果模型把填充位置当成了有效 token 来学，损失死活降不下去。

4.2 填充策略——不只是简单补零

填充其实有几种策略，我按实用程度排个序：

策略	做法	适用场景
右填充	在序列末尾补 pad	大多数 NLP 任务，简单直观
左填充	在序列开头补 pad	某些生成任务，如 GPT 风格
动态批次	按长度排序后分组填充	训练效率优先，减少无效计算

我个人最常用的是动态批次。什么意思呢？就是先把所有样本按长度排序，然后把长度相近的放在同一个批次里。这样每个批次里的序列长度差异不大，填充量就少。

举个例子：

def dynamic_batch_collate(batch, bucket_size=32):
    # 按输入长度排序
    batch.sort(key=lambda x: len(x[0]))
    
    # 分成多个 bucket
    buckets = []
    for i in range(0, len(batch), bucket_size):
        buckets.append(batch[i:i+bucket_size])
    
    # 每个 bucket 内部做填充
    padded_batches = []
    for bucket in buckets:
        max_len = max(len(item[0]) for item in bucket)
        inputs, targets = [], []
        for inp, tgt in bucket:
            pad_len = max_len - len(inp)
            inputs.append(inp + [0] * pad_len)
            targets.append(tgt + [0] * pad_len)
        padded_batches.append((torch.tensor(inputs), torch.tensor(targets)))
    
    return padded_batches

这样做的好处很明显——训练速度能提升 30%-50%。我在做 BERT 预训练时，用动态批次把原本 7 天的训练时间压缩到了 4 天半。

4.3 掩码——告诉模型哪里是填充

填充完了，模型怎么知道哪些位置是真实数据，哪些是填充的？

答案就是注意力掩码。

Transformer 的自注意力机制会计算每个位置对其他位置的注意力权重。如果不加掩码，模型会把填充位置的向量也纳入计算，相当于引入了噪声。

掩码矩阵长这样：

def create_padding_mask(seq, pad_token_id=0):
    # seq shape: [batch_size, seq_len]
    mask = (seq != pad_token_id).unsqueeze(1).unsqueeze(2)
    # mask shape: [batch_size, 1, 1, seq_len]
    return mask

# 举个例子
seq = torch.tensor([[3, 5, 7, 0, 0], [2, 4, 0, 0, 0]])
mask = create_padding_mask(seq)
print(mask)
# 输出:
# tensor([[[[ True,  True,  True, False, False]]],
#         [[[ True,  True, False, False, False]]]])

为什么是 4 维？因为 Transformer 的注意力分数形状是 [batch_size, num_heads, seq_len, seq_len]，掩码需要能广播到这个形状。

4.4 因果掩码——防止看到未来

除了填充掩码，还有一类掩码叫因果掩码（也叫自回归掩码）。

你想想看，在生成任务里，模型预测第 3 个词时，不能看到第 4 个词的内容。否则就是作弊了。

因果掩码是一个上三角矩阵：

def create_causal_mask(seq_len):
    # 上三角矩阵，右上角为 False
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1) == 0
    return mask

# 长度为 4 的因果掩码
mask = create_causal_mask(4)
print(mask)
# 输出:
# tensor([[ True, False, False, False],
#         [ True,  True, False, False],
#         [ True,  True,  True, False],
#         [ True,  True,  True,  True]])

在实际工程中，我们通常把填充掩码和因果掩码合并成一个：

def combine_masks(padding_mask, causal_mask):
    # padding_mask: [batch, 1, 1, seq_len]
    # causal_mask: [1, 1, seq_len, seq_len]
    return padding_mask & causal_mask

嗯，这里有个细节：padding_mask 是 4 维的，causal_mask 是 4 维的，两者做 & 运算时 PyTorch 会自动广播。

4.5 完整的数据加载器实现

把上面这些整合到一起，就是一个完整的数据加载器：

class TransformerDataLoader:
    def __init__(self, dataset, batch_size=32, pad_token_id=0, shuffle=True):
        self.dataset = dataset
        self.batch_size = batch_size
        self.pad_token_id = pad_token_id
        self.shuffle = shuffle
        
    def __iter__(self):
        indices = list(range(len(self.dataset)))
        if self.shuffle:
            import random
            random.shuffle(indices)
        
        # 按长度排序（动态批次）
        indices.sort(key=lambda i: len(self.dataset[i][0]))
        
        for i in range(0, len(indices), self.batch_size):
            batch_indices = indices[i:i+self.batch_size]
            batch = [self.dataset[idx] for idx in batch_indices]
            
            # 填充
            max_len = max(len(item[0]) for item in batch)
            inputs, targets = [], []
            for inp, tgt in batch:
                pad_len = max_len - len(inp)
                inputs.append(inp + [self.pad_token_id] * pad_len)
                targets.append(tgt + [self.pad_token_id] * pad_len)
            
            inputs = torch.tensor(inputs)
            targets = torch.tensor(targets)
            
            # 生成掩码
            padding_mask = create_padding_mask(inputs, self.pad_token_id)
            causal_mask = create_causal_mask(max_len)
            combined_mask = combine_masks(padding_mask, causal_mask)
            
            yield inputs, targets, combined_mask

好了，数据组件到这里就讲完了。下一章我们开始搭建 Transformer 的核心——注意力机制。说实话，这部分才是真正的硬核内容，我当年啃了整整两周才完全搞明白。不过别担心，我会用最直观的方式讲给你听。