这个函数 `get_attn_pad_mask` 的作用是在序列到序列模型（如Transformer）中生成注意力机制使用的 **pad mask** 矩阵，用于屏蔽填充位置的信息。以下是详细解释： ### 参数说明： 1. **seq_q**: 查询序列 (Query Sequence)，形状为 `(batch_size, len_q)`。 - 表示当前批次所有查询的 token 序列。 2. **seq_k**: 键序列 (Key Sequence)，形状为 `(batch_size, len_k)`。 - 表示当前批次所有的键值对 token 序列。 3. **pad_index**: 填充符的位置索引，在大多数任务中通常是 `<PAD>` 标记对应的词表索引。 --- ### 函数功能解析： #### 第一步：获取输入维度 ```python batch_size, len_q = seq_q.size() batch_size, len_k = seq_k.size() ``` 这里提取了批量大小 (`batch_size`) 和两个序列长度 (`len_q`, `len_k`)。 #### 第二步：创建掩码矩阵 ```python pad_attn_mask = seq_k.data.eq(pad_index).unsqueeze(1) ``` 这一步的核心操作是找到 `seq_k` 中等于 `pad_index` 的位置，并将其标记为 `True`（表示需要屏蔽的地方），其他地方为 `False`。然后通过 `.unsqueeze(1)` 将其扩展成三维张量 `[batch_size, 1, len_k]` 形状。 #### 第三步：将布尔型转换为整数型 ```python pad_attn_mask = torch.as_tensor(pad_attn_mask, dtype=torch.int) ``` 为了方便后续计算，通常会把布尔类型的掩码转为整数类型（0 或 1）。例如： - 如果某个位置对应的是 `<PAD>` 符号，则该位置赋值为 `1`； - 其他正常字符则保持为 `0`。 #### 第四步：扩展掩码尺寸 ```python return pad_attn_mask.expand(batch_size, len_q, len_k) ``` 最后使用 `.expand()` 方法将掩码从 `[batch_size, 1, len_k]` 扩展至 `[batch_size, len_q, len_k]`，使得每个查询都可以独立地应用相同的 key 掩码规则。 --- ### 返回结果的意义 返回的结果是一个三维张量 `[batch_size, len_q, len_k]`，其中： - 每个元素都指示了是否应该忽略掉 `key` 对应的位置信息（如果某位置是 `1` 则会被屏蔽）。 - 这种设计可以避免 Transformer 在处理变长句子时受到无意义填充内容的影响，提升性能和稳定性。 --- ### 示例代码运行过程演示： 假设我们有如下数据作为输入： ```python import torch # 定义输入序列及 pad_index seq_q = torch.tensor([[1, 2, 0], [4, 5, 6]]) # shape: [2, 3] seq_k = torch.tensor([[7, 8, 0], [9, 10, 11]]) # shape: [2, 3] pad_index = 0 mask = get_attn_pad_mask(seq_q, seq_k, pad_index) print(mask.shape) # 输出应该是 [2, 3, 3] print(mask) ``` 得到的输出可能是这样的： ```plaintext tensor([[[0, 0, 1], [0, 0, 1], [0, 0, 1]], [[0, 0, 0], [0, 0, 0], [0, 0, 0]]]) ``` 可以看到第一行最后一列为全 `1`，因为第三个单词已经是 padding 数据；而第二句没有多余的填充部分所以全是零。 ---