pytorch也自己实现了transformer的模型，不同于huggingface或者其他地方，pytorch的mask参数要更难理解一些（即便是有文档的情况下），这里做一些补充和说明。（顺带提一句，这里的transformer是需要自己实现position embedding的，别乐呵乐呵的就直接去跑数据了）
代码语言：javascript复制>>> transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
>>> src = torch.rand((10, 32, 512))
>>> tgt = torch.rand((20, 32, 512))
>>> out = transformer_model(src, tgt) # 没有实现position embedding ，也需要自己实现mask机制。否则不是你想象的transformer首先看一下官网的参数
- src – the sequence to the encoder (required).
- tgt – the sequence to the decoder (required).
- src_mask – the additive mask for the src sequence (optional).
- tgt_mask – the additive mask for the tgt sequence (optional).
- memory_mask – the additive mask for the encoder output (optional).
- src_key_padding_mask – the ByteTensor mask for src keys per batch (optional).
- tgt_key_padding_mask – the ByteTensor mask for tgt keys per batch (optional).
- memory_key_padding_mask – the ByteTensor mask for memory keys per batch (optional).
这里面最大的区别就是*mask_和*_key_padding_mask,_至于*是src还是tgt，memory，这不重要，模块出现在encoder，就是src，出现在decoder，就是tgt，decoder每个block的第二层和encoder做cross attention的时候，就是memory。
*mask 对应的API是attn_mask，*_key_padding_mask对应的API是key_padding_mask
我们看看torch/nn/modules/activation.py当中MultiheadAttention模块对于这2个API的解释：
代码语言：javascript复制def forward(self, query, key, value, key_padding_mask=None,
need_weights=True, attn_mask=None):
# type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
r"""
Args:
query, key, value: map a query and a set of key-value pairs to an output.
See "Attention Is All You Need" for more details.
key_padding_mask: if provided, specified padding elements in the key will
be ignored by the attention. When given a binary mask and a value is True,
the corresponding value on the attention layer will be ignored. When given
a byte mask and a value is non-zero, the corresponding value on the attention
layer will be ignored
need_weights: output attn_output_weights.
attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
the batches while a 3D mask allows to specify a different mask for the entries of each batch.
Shape:
- Inputs:
- query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
the embedding dimension.
- key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
the embedding dimension.
- value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
the embedding dimension.
- key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
If a ByteTensor is provided, the non-zero positions will be ignored while the position
with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
- attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
is provided, it will be added to the attention weight.
- Outputs:
- attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
E is the embedding dimension.
- attn_output_weights: :math:`(N, L, S)` where N is the batch size,
L is the target sequence length, S is the source sequence length.
"""
- key_padding_mask：用来遮蔽<PAD>以避免pad token的embedding输入。形状要求：（N,S）
- attn_mask：2维或者3维的矩阵。用来避免指定位置的embedding输入。2维矩阵形状要求：（L, S）；也支持3维矩阵输入，形状要求：（N*num_heads, L, S）
其中，N是batch size的大小，L是目标序列的长度(the target sequence length)，S是源序列的长度(the source sequence length)。这个模块会出现在上图的3个橙色区域，所以the target sequence 并不一定就是指decoder输入的序列，the source sequence 也不一定就是encoder输入的序列。
更准确的理解是，target sequence代表多头attention当中q（查询）的序列，source sequence代表k（键值）和v（值）的序列。例如，当decoder在做self-attention的时候，target sequence和source sequence都是它本身，所以此时L=S，都是decoder编码的序列长度。
key_padding_mask的作用
这里举一个简单的例子：
现在有一个batch，batch_size = 3，长度为4，token表现形式如下：
代码语言：javascript复制[
[‘a’,'b','c','<PAD>'],
[‘a’,'b','c','d'],
[‘a’,'b','<PAD>','<PAD>']
]
现在假设你要对其进行self-attention的计算（可以在encoder，也可以在decoder），那么以第三行数据为例，‘a’在做qkv计算的时候，会看到'b','<PAD>','<PAD>'，但是我们不希望‘a’看到'<PAD>'，因为他们本身毫无意义，所以，需要key_padding_mask遮住他们。
key_padding_mask的形状大小为（N,S），对应这个例子，key_padding_mask为以下形式，key_padding_mask.shape = （3,4）：
代码语言：javascript复制[
[False, False, False, True],
[False, False, False, False],
[False, False, True, True]
]
值得说明的是，key_padding_mask本质上是遮住key这个位置的值（置0），但是<PAD> token本身，也是会做qkv的计算的，以第三行数据的第三个位置为例，它的q是<PAD>的embedding，k和v分别各是第一个的‘a’和第二个的‘b’，它也会输出一个embedding。
所以你的模型训练在transformer最后的output计算loss的时候，还需要指定ignoreindex=pad_index。以第三行数据为例，它的监督信号是[3205,1890,0,0]，pad_index=0 。如此一来，即便位于<PAD>的transformer会疯狂的和有意义的position做qkv，也会输出embedding，但是我们不算它的loss，任凭它各种作妖。
attn_mask的作用
一开始看到有2个mask参数的时候，我也是一脸懵逼的，并且他们的shape居然要求还不一样。attn_mask到底用在什么地方呢？
decoder在做self-attention的时候，每一个位置不同于encoder，他是只能看到上文的信息的。key_padding_mask的shape为(batch_size, source_length)，这意味着每个位置的query，他所看到的画面经过key_padding_mask后都是一样的（尽管他能做到batch的每一行数据mask的不一样），这不能满足如下模块的需求：
decoder的mask 多头注意力模块
这里需要的mask如下：
黄色是看得到的部分，紫色是看不到的部分，不同位置需要mask的部分是不一样的
而pytorch的nn.Transformer已经有了帮我们实现的函数：
代码语言：javascript复制 def generate_square_subsequent_mask(self, sz: int) -> Tensor:
r"""Generate a square mask for the sequence. The masked positions are filled with float('-inf').
Unmasked positions are filled with float(0.0).
"""
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
还是上面那个例子，以第一行数据['a','b','c','<PAD>'],为例（假设我们在用decoder做生成，研究block 的第一层layer 也就是self-attention），此时：
- 'a'可以看到'a'
- 'b'可以看到'a','b'
- 'c'可以看到'a','b','c'
- '<PAD>'理论上不应该看到什么，但是只要它头顶的监督信号是ignore_index，那就没有关系，所以让他看到'a','b','c','<PAD>'
回想一下attn_mask的形状要求，2维的时候是（L,S），3维的时候是（N*num_heads, L, S）。此时，由于qkv都是同一个序列（decoder底下的序列）所以L=S；又因为对于batch每一行数据来说，他们的mask机制都是一样的，即第i个位置的值，都只能看到上文的信息，所以我们的attn_mask用二维的就行，内部实现的时候会把mask矩阵广播到batch每一行数据中：
一般而言，除非你需要魔改transformer，例如让不同的头看不同的信息，否则二维的矩阵足够使用了。
什么时候用key_padding_mask，什么时候用attn_mask？
个人感觉最好是按照上面的约定的习惯来用，实际上，2个mask共同作用于同一个模型，非要用attn_mask代替key_padding_mask把<PAD>遮住，行不行？当然可以，只不过这只会增加你的工作量。
说个正事哈
