PyTorch序列数据编码中避免填充（Padding）影响的策略

import torch

# 示例数据
bs = 2  # 批次大小
sl = 5  # 序列长度 (包含填充)
n = 10  # 特征维度

# 假设这是模型输出的序列嵌入 (bs, sl, n)
# 为了演示，我们手动创建一个带有填充值的张量
embeddings = torch.randn(bs, sl, n)
# 模拟填充：例如，第一个序列真实长度为3，第二个序列真实长度为4
# 填充部分我们将其设置为0，以更清晰地看到掩码的作用
embeddings[0, 3:] = 0.0
embeddings[1, 4:] = 0.0

print("原始嵌入 (部分填充为0):\n", embeddings)

# 对应的二进制填充掩码 (bs, sl)
# 1 表示非填充，0 表示填充
padding_mask = torch.tensor([
    [1, 1, 1, 0, 0],  # 第一个序列的真实长度是3
    [1, 1, 1, 1, 0]   # 第二个序列的真实长度是4
], dtype=torch.float32)

print("\n填充掩码:\n", padding_mask)

# 1. 扩展掩码维度以匹配嵌入的特征维度
# padding_mask.unsqueeze(-1) 将形状从 (bs, sl) 变为 (bs, sl, 1)
# 这样就可以与 (bs, sl, n) 的 embeddings 进行广播乘法
expanded_mask = padding_mask.unsqueeze(-1)
print("\n扩展后的掩码形状:", expanded_mask.shape)

# 2. 将嵌入与扩展后的掩码相乘
# 这一步会将填充位置的嵌入值变为0，非填充位置保持不变
masked_embeddings = embeddings * expanded_mask
print("\n应用掩码后的嵌入 (填充部分变为0):\n", masked_embeddings)

# 3. 对掩码后的嵌入在序列长度维度上求和
# sum(1) 会将 (bs, sl, n) 变为 (bs, n)
sum_masked_embeddings = masked_embeddings.sum(1)
print("\n求和后的嵌入:\n", sum_masked_embeddings)

# 4. 计算每个序列中非填充元素的数量
# padding_mask.sum(-1) 将形状从 (bs, sl) 变为 (bs,)
# 然后 unsqueeze(-1) 变为 (bs, 1)，以便进行广播除法
non_padding_counts = padding_mask.sum(-1).unsqueeze(-1)
# 使用 torch.clamp 避免除以零的情况，当序列完全由填充组成时
non_padding_counts_clamped = torch.clamp(non_padding_counts, min=1e-9)
print("\n非填充元素数量:\n", non_padding_counts_clamped)

# 5. 计算均值嵌入
mean_embeddings = sum_masked_embeddings / non_padding_counts_clamped
print("\n最终的均值嵌入 (形状: {}, 不含填充):\n".format(mean_embeddings.shape), mean_embeddings)

# 验证结果：手动计算第一个序列的均值
# 真实数据点：embeddings[0, 0], embeddings[0, 1], embeddings[0, 2]
# expected_mean_0 = (embeddings[0, 0] + embeddings[0, 1] + embeddings[0, 2]) / 3
# print("\n手动计算第一个序列的均值:\n", expected_mean_0)
# print("与模型计算结果的差异 (第一个序列):", (mean_embeddings[0] - expected_mean_0).abs().sum())