第3节 网络结构与自注意力实现#
在通过前面几部分内容详细介绍完Transformer网络结构的原理后，接下来就让我们来看一看如何借用Pytorch框架来实现 MultiHeadAttention 这一结构。同时，需要说明的一点是，下面所有的实现代码都是直接从 Pytorch 1.4 版本中 torch.nn.Transformer 模块里摘取出来的简略版，目的就是为了让大家对于整个实现过程有一个清晰的认识。并且为了使得大家在阅读完以下内容后也能够对Pytorch中的相关模块有一定的了解，所以下面的代码在变量名方面也与Pytorch保持了一致。
代码仓库见：https://github.com/mlwithme/TransformerTranslation
3.1多层Transformer#
在第2部分中，我们详细介绍了单层Transformer网络结构中的各个组成部分。尽管多层Transformer就是在此基础上堆叠而来，不过还是有必要在这里稍微提及一下。
如图3-1所示便是一个单层Transformer网络结构图，左边是编码器右边是解码器。而多层的Transformer网络就是在两边分别堆叠了多个编码器和解码器的网络模型，如图3-2所示。
如图3-2所示便是一个多层的Transformer网络结构图（原论文中采用了6个编码器和6个解码器），其中的每一个Encoder都是图3-1中左边所示的网络结构（Decoder同理）。可以发现，它真的就是图3-1堆叠后的形式。不过需要注意的是其整个解码过程。
在多层Transformer中，多层编码器先对输入序列进行编码，然后得到最后一个Encoder的输出Memory；解码器先通过Masked Multi-Head Attention对输入序列进行编码，然后将输出结果同Memory通过Encoder-Decoder Attention后得到第1层解码器的输出；接着再将第1层Decoder的输出通过Masked Multi-Head Attention进行编码，最后再将编码后的结果同Memory通过Encoder-Decoder Attention后得到第2层解码器的输出，以此类推得到最后一个Decoder的输出。
值得注意的是，在多层Transformer的解码过程中，每一个Decoder在Encoder-Decoder Attention中所使用的Memory均是同一个。
3.2 Transformer中的掩码#
由于在实现多头注意力时需要考虑到各种情况下的掩码，因此在这里需要先对这部分内容进行介绍。在Transformer中，主要有两个地方会用到掩码这一机制。第1个地方就是在2.2.4节中介绍到的Attention Mask，只用于在训练过程中解码阶段时掩盖掉当前时刻之后的信息；第2个地方便是对一个batch中不同长度的序列在Padding到相同长度后，对Padding部分的信息进行掩盖，这在编码器和解码器中均会用到。下面分别就这两种情况进行介绍。
3.2.1 Attention Mask#
如图3-3所示，在训练过程中对于每一个样本来说都需要这样一个对称矩阵来掩盖掉当前时刻之后所有位置的信息。
从图3-3可以看出，这个注意力掩码矩阵的形状为 [tgt_len,tgt_len]，其具体Mask原理在2.2.4节中已经介绍过了，这里就不再赘述。在后续实现过程中，我们将通过 generate_square_subsequent_mask 方法来生成这样一个矩阵。同时，在后续多头注意力机制实现中，将通过attn_mask这一变量名来指代这个矩阵。
3.2.2 Padding Mask#
在Transformer中，使用到掩码的第2个地方便是Padding Mask。由于在网络的训练过程中同一个batch会包含有多个文本序列，而不同的序列长度并不一致。因此在数据集生成时，就需要将同一个batch中的序列Padding到相同的长度。但是，这样就会导致在注意力的计算过程中会考虑到Padding位置上的信息。
如图3-4所示，P表示Padding位置，右边的矩阵表示注意力权重矩阵。从图中可以看到，此时注意力权重对于Padding位置上的信息也会加以考虑。因此在Transformer中，作者通过在生成训练集的过程中记录下每个样本Padding的实际位置；然后再将注意力权重矩阵中对应位置的权重替换成负无穷，经softmax操作后对应Padding位置上的权重就变成了0，从而达到了忽略Padding位置信息的目的。这种做法也是Encoder-Decoder网络结构中通用的一种办法。
如图3-5所示，对于 "我 是 谁 P P" 这个序列来说，前3个字符是正常的，后2个字符是Padding后的结果。因此，其Mask向量便为 [True, True, True, False, False] 。通过这个Mask向量可知，需要将权重矩阵的最后两列替换成负无穷，在后续我们会通过 torch.masked_fill 这个方法来完成这一步，并且在实现时将使用 key_padding_mask 来指代这一向量。
3.2.3 行Padding与列Padding#
在3.2.2节中我们介绍到可以通过 mask fill 操作来忽略掉注意力权重矩阵中padding部分的信息，即图3-5中的示例。但此时有好事者发出疑问，图3-5中注意力权重的最后两行为什么没有进行mask操作？这两行对应的不也是padding部分的信息吗？此时，问题就来了，为什么不对注意力权重矩阵中的最后两行也进行mask处理（如图3-6右所示），它代表的不也是padding部分的信息吗？
这里之所以不用对最后两行的padding进行mask处理，是因为权重矩阵与V作用后得到的结果Memory在与Decoder中Multi-Head Attention交互时，最后两行padding部分的内容在那里进行了mask处理。下面，我们就用图示来一步一步还原整个过程，以便大家能够清楚地理解。
假如图3-6中左边的权重矩阵与V作用后的结果如图3-7所示，此时可以知道Z中最后两行的信息本质上也是属于padding部分的内容，按道理来说如果权重矩阵的最后两行全为0 ，那么Z中最后两行的结果也将全部为0而被忽略掉。
进一步，假定解码器的输入经过图3-1中Masked Multi-Head Attention处理后的输出为tgt(Q)，与编码器的输出Memory（图3-7中Z经过线性变换后的结果）作用后得到权重矩阵如图3-8最右边所示。按道理，如果图3-7中Z的最后两行均为0，那么图3-8中权重矩阵的最后两列将会全部为0。
接着，图3-8中的注意力矩阵将会进行类似图3-5中的mask操作，而此时encoder-decoder交互部分的maks向量就是图3-5中的mask向量。因此，整个mask过程如图3-9所示。
如图3-9所示，图3-8中的注意力矩阵在经过mask操作后，注意力矩阵的最后两列全部变成了0，从而也就达到了后面可以忽略掉图3-7里Z中最后两行信息的目的，如图3-10所示。
如图3-10所示，由于此时注意力矩阵的最后两列均为0，那么其在与encoder输出部分memory作用时，memory中最后两行padding部分的信息自然就得到了忽略。此时可以看出，即使在encoder里图3-7中权重矩阵的最后两行没有进行mask，但是在encoder-decoder交互时的mask操作中却得到了体现。不过，假如图3-7中权重矩阵的最后两行也进行了mask，那又会是一个什么样的结果呢？
如图3-11所示，假如此时权重矩阵的最后两行也进行了padding处理，那么此时encoder单头注意力的输出结果将形如Z的形式，即padding部分的信息全为0。进一步，图3-11中Z经线性变换得到memory(K)后，与解码器的输入经过Masked Multi-Head Attention处理后的输出tgt(Q)作用后的结果将如图3-12所示。
如图3-12所示，由于此时memory中最后两行的信息也进行了mask处理，所以计算得到的注意力权重矩阵最后两列便直接变成了0，从而也直接得到了图3-9中mask操作后的结果，后续便再是图3-10所示的过程。
到此可以得出的结论便是，在encoder中无论是否要在行上注意力权重矩阵进行mask操作，都不会影响模型最后的输出结果。因为即使在encoder中注意力权重矩阵行上的padding信息没有被mask，但是在encoder-decoder交互部分的注意力权重矩阵计算过程中也会对这部分信息进行mask处理。
不过最后的最后，还要再提出一个类似的问题，在解码器的Masked Multi-Head Attention中，注意力权重矩阵同样也只是在列上进行了mask处理，那为什么同样没有在行上进行mask操作呢？或者说后续阶段哪个操作等价于行上的mask操作呢？如果想不明白可以公众号后台回复“行padding”获取答案。
到此，对于Transformer中所要用到Mask的地方就介绍完了，下面正式来看如何实现多头注意力机制。
3.3 实现多头注意力机制#
3.3.1 多头注意力机制#
根据前面的介绍可以知道，多头注意力机制中最为重要的就是自注意力机制，也就是需要前计算得到Q、K和V，如图3-13所示。
然后再根据Q、K、V来计算得到最终的注意力编码，如图3-14所示。
同时，为了避免单个自注意力机制计算得到的注意力权重过度集中于当前编码位置自己所在的位置（同时更应该关注于其它位置），所以作者在论文中提到通过采用多头注意力机制来解决这一问题，如图3-15所示。
3.3.2 定义类 MyMultiHeadAttention#
综上所述，我们可以给出类 MyMultiHeadAttentiond 的定义为
1 class MyMultiheadAttention(nn.Module):
2 def __init__(self,embed_dim, num_heads, dropout=0.,bias=True):
3 super(MyMultiheadAttention, self).__init__()
4 """
5 :param embed_dim: 词嵌入的维度，也就是前面的d_model参数,论文中的默认值为512
6 :param num_heads:多头注意力中多头的数量,也就是前面的nhead参数,论文默认值为 8
7 :param bias: 最后对多头的注意力（组合）输出进行线性变换时，是否使用偏置
8 """
9 self.embed_dim = embed_dim # 前面的d_model参数
10 self.head_dim = embed_dim //num_heads # head_dim 指的就是d_k,d_v
11 self.kdim = self.head_dim
12 self.vdim = self.head_dim
13 self.num_heads = num_heads # 多头个数
14 self.dropout = dropout
15 assert self.head_dim * num_heads == self.embed_dim,
16 "embed_dim 除以 num_heads必须为整数"
17 # 上面的限制条件就是论文中的 d_k = d_v = d_model/n_head 条件
18 self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
19 # embed_dim = kdim * num_heads
20 # 这里第二个维度之所以是embed_dim，实际上这里是同时初始化了num_heads个W_q 堆叠起来的, 也就是num_heads个头
21 self.k_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
22 # W_k, embed_dim = kdim * num_heads
23 self.v_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
24 # W_v, embed_dim = vdim * num_heads
25 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
26 # 最后将所有的Z组合起来的时候，也是一次性完成， embed_dim = vdim * num_heads在上述代码中，embed_dim表示模型的维度（图3-15中的d_m）；num_heads表示多头的个数；bias表示是否在多头线性组合时使用偏置。同时，为了使得实现代码更加高效，所以Pytorch在实现的时候是多个头注意力机制一起进行的计算，也就上面代码的第17~22行，分别用来初始化了多个头的权重值（这一过程从图3-15也可以看出）。当多头注意力机制计算完成后，将会得到一个形状为[src_len,embed_dim]的矩阵，也就是图3-15中多个 水平堆叠后的结果。因此，第24行代码将会初始化一个线性层来对这一结果进行一个线性变换。
3.3.3 定义前向传播过程#
在定义完初始化函数后，便可以定义如下所示的多头注意力前向传播的过程
1 def forward(self, query,key,value,attn_mask = None, key_padding_mask=None):
2 """
3 在论文中，编码时query, key, value 都是同一个输入，
4 解码时 输入的部分也都是同一个输入，
5 解码和编码交互时 key,value指的是 memory, query指的是tgt
6 :param query: [tgt_len,batch_size,embed_dim],tgt_len 表示目标序列长度
7 :param key: [src_len, batch_size,embed_dim], src_len表示源序列长度
8 :param value: [src_len, batch_size,embed_dim], src_len表示源序列长度
9 :param attn_mask:[tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len] 一般只在解码时使用，为了并行一次喂入所有解码部分的输入，
10 所以要用mask来进行掩盖当前时刻之后的位置信息
11 :param key_padding_mask: [batch_size, src_len], src_len 表示源序列长度
12 :return:
13 attn_output: [tgt_len, batch_size, embed_dim]
14 attn_output_weights: # [batch_size, tgt_len, src_len]
15 """
16 return multi_head_attention_forward(query, key,value,self.num_heads,
17 self.dropout,self.out_proj.weight,
18 self.out_proj.bias, training=self.training,
19 key_padding_mask=key_padding_mask,
20 q_proj_weight=self.q_proj_weight,
21 k_proj_weight=self.k_proj_weight,
22 v_proj_weight=self.v_proj_weight,
23 attn_mask=attn_mask)在上述代码中，query、key、value指的并不是图3-13中的Q、K和V，而是没有经过线性变换前的输入。例如在编码时三者均是指原始输入序列src；在解码时的Mask Multi-Head Attention中三者指的均是目标输入序列tgt；在解码时的Encoder-Decoder Attention中三者分别指的是Mask Multi-Head Attention的输出、Memory和Memory。key_padding_mask指的是编码或解码部分，输入序列的Padding情况，形状为 [batch_size,src_len] 或者[batch_size,tgt_len];attn_mask指的就是注意力掩码矩阵，形状为 [tgt_len,src_len]，它只会在解码时使用。
注意，在上面的这些维度中，tgt_len本质上指的其实是query_len；src_len本质上指的是key_len。只是在不同情况下两者可能会是一样，也可能会是不一样。
3.3.4 多头注意力计算过程#
在定义完类 MyMultiHeadAttentiond 后，就需要定义出多头注意力的实际计算过程。由于这部分代码较长，所以就分层次进行介绍。
1 def multi_head_attention_forward(
2 query, # [tgt_len,batch_size, embed_dim]
3 key, # [src_len, batch_size, embed_dim]
4 value, # [src_len, batch_size, embed_dim]
5 num_heads,
6 dropout_p,
7 out_proj_weight, # [embed_dim = vdim * num_heads, embed_dim]
8 out_proj_bias,
9 training=True,
10 key_padding_mask=None, # [batch_size,src_len/tgt_len]
11 q_proj_weight=None, # [embed_dim,kdim * num_heads]
12 k_proj_weight=None, # [embed_dim, kdim * num_heads]
13 v_proj_weight=None, # [embed_dim, vdim * num_heads]
14 attn_mask=None, # [tgt_len,src_len]
15 ):
16 # 第一阶段： 计算得到Q、K、V
17 q = F.linear(query, q_proj_weight)
18 # [tgt_len,batch_size,embed_dim] x [embed_dim,kdim * num_heads]
19 # = [tgt_len,batch_size,kdim * num_heads]
20 k = F.linear(key, k_proj_weight)
21 # [src_len, batch_size,embed_dim] x [embed_dim,kdim * num_heads]
22 # = [src_len,batch_size,kdim * num_heads]
23 v = F.linear(value, v_proj_weight)
24 # [src_len, batch_size,embed_dim] x [embed_dim,vdim * num_heads]
25 # = [src_len,batch_size,vdim * num_heads]在上述代码中，第17~23行所做的就是根据输入进行线性变换得到图3-13中的Q、K和V。
1 # 第二阶段： 缩放，以及attn_mask维度判断
2 tgt_len, bsz, embed_dim = query.size()# [tgt_len,batch_size,embed_dim]
3 src_len = key.size(0)
4 head_dim = embed_dim // num_heads # num_heads * head_dim = embed_dim
5 scaling = float(head_dim) ** -0.5
6 q = q * scaling # [query_len,batch_size,kdim * num_heads]
7
8 if attn_mask is not None:
9 # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
10 if attn_mask.dim() == 2:
11 attn_mask =attn_mask.unsqueeze(0)#[1,tgt_len,src_len]扩充维度
12 if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
13 raise RuntimeError('The size of the 2D attn_mask is not correct.')
14 elif attn_mask.dim() == 3:
15 if list(attn_mask.size()) != [bsz * num_heads, query.size(0),key.size(0)]:
16 raise RuntimeError('The size of the 3D attn_mask is not correct.')
17 # 现在 atten_mask 的维度就变成了3D接着，在上述代码中第5~6行所完成的就是图3-14中的缩放过程；第8~16行用来判断或修改attn_mask的维度，当然这几行代码只会在解码器中的Masked Multi-Head Attention中用到。
1 # 第三阶段： 计算得到注意力权重矩阵
2 q =q.contiguous().view(tgt_len, bsz*num_heads,head_dim).transpose(0, 1)
3 # [batch_size * num_heads,tgt_len,kdim]
4 # 因为前面是num_heads个头一起参与的计算，所以这里要进行一下变形，以便于 后面计算。且同时交换了0，1两个维度
5 k = k.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
6 #[batch_size * num_heads,src_len,kdim]
7 v = v.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
8 #[batch_size * num_heads,src_len,vdim]
9 attn_output_weights = torch.bmm(q, k.transpose(1, 2))
10 # [batch_size * num_heads,tgt_len,kdim] x [batch_size * num_heads, kdim, src_len] = [batch_size*num_heads, tgt_len, src_len]
11 # 这就num_heads个QK相乘后的注意力矩阵继续，在上述代码中第1~7行所做的就是交换Q、K、V中的维度，以便于多个样本同时进行计算；第9行代码便是用来计算注意力权重矩阵；其中上contiguous()方法是将变量放到一块连续的物理内存中；bmm的作用是用来计算两个三维矩阵的乘法操作[4]。
需要提示的是，大家在看代码的时候，最好是仔细观察一下各个变量维度的变化过程，我们也在每次运算后进行了批注。
1 # 第四阶段： 进行相关掩码操作
2 if attn_mask is not None:
3 attn_output_weights += attn_mask #[batch_size*num_heads,tgt_len,src_len]
4 if key_padding_mask is not None:
5 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
6 # 变成 [batch_size, num_heads, tgt_len, src_len]的形状
7 attn_output_weights = attn_output_weights.masked_fill(
8 key_padding_mask.unsqueeze(1).unsqueeze(2),float('-inf'))
9 # 扩展维度，从[batch_size,src_len]变成[batch_size,1,1,src_len]
10 attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len,src_len)
11 # [batch_size * num_heads, tgt_len, src_len]进一步，在上述代码中第2~3行便是用来执行图3-3中的步骤；第4~8行便是用来执行图3-5中的步骤，同时还进行了维度扩充。
1 attn_output_weights = F.softmax(attn_output_weights, dim=-1)
2 # [batch_size * num_heads, tgt_len, src_len]
3 attn_output_weights=F.dropout(attn_output_weights,p=dropout_p, training=training)
4 attn_output = torch.bmm(attn_output_weights, v)
5 # Z = [batch_size*num_heads,tgt_len, src_len]@[batch_size * num_heads, src_len,vdim] = [batch_size * num_heads,tgt_len,vdim]
6
7 # 这就num_heads个Attention(Q,K,V)结果
8 attn_output = attn_output.transpose(0,1).contiguous().view(tgt_len,bsz, embed_dim)
9 # 先transpose成 [tgt_len, batch_size* num_heads ,kdim]
10 # 再view成 [tgt_len,batch_size,num_heads*kdim]
11 attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len,src_len)
12
13 Z = F.linear(attn_output, out_proj_weight, out_proj_bias)
14 # 这里就是多个z 线性组合成Z [tgt_len,batch_size,embed_dim]
15 return Z, attn_output_weights.sum(dim=1) / num_heads #将num_heads个注意
# 力权重矩阵按对应维度取平均最后，在上述代码中第1~4行便是用来对权重矩阵进行归一化操作，以及计算得到多头注意力机制的输出；第13行代码便是用来对多个注意力的输出结果进行线性组合；第15行代码用来返回线性组合后的结果，以及多个注意力权重矩阵的平均值。
3.3.5 示例代码#
在实现完类MyMultiHeadAttention的全部代码后，便可以通过类似如下的方式进行使用。
1 if __name__ == '__main__':
2 src_len = 5
3 batch_size = 2
4 dmodel = 32
5 num_head = 1
6 src = torch.rand((src_len, batch_size, dmodel))
7 # shape: [src_len, batch_size, embed_dim]
8 src_key_padding_mask = torch.tensor([
9 [True, True, True, False, False],
10 [True, True, True, True, False]])# shape:[src_len,src_len]
11 my_mh = MyMultiheadAttention(embed_dim=dmodel, num_heads=num_head)
12 r = my_mh(src,src,src,key_padding_mask = src_key_padding_mask)在上述代码中，第6~10行其实也就是Encoder中多头注意力机制的实现过程。同时，在计算过程中还可以打印出各个变量的维度变化信息：
1 进入多头注意力计算
2 多头num_heads = 1, d_model=32, d_k = d_v = d_model/num_heads=32
3 query 的shape([tgt_len, batch_size, embed_dim]):torch.Size([5, 2, 32])
4 W_q 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
5 Q的shape([tgt_len, batch_size,kdim * num_heads]):torch.Size([5,2,32])
6 ----------------------------------------------------------------------
7 key 的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
8 W_k 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32])
9 K 的shape([src_len,batch_size,kdim * num_heads]):torch.Size([5,2,32])
10 ----------------------------------------------------------------------
11 value的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32])
12 W_v 的shape([embed_dim,vdim * num_heads]):torch.Size([32, 32])
13 V 的shape([src_len,batch_size,vdim * num_heads]):torch.Size([5,2,32])
14 ----------------------------------------------------------------------
15 ***** 注意，这里的W_q, W_k, W_v是多个head同时进行计算的. 因此Q,K,V分别也是包含了多个head的 q,k,v堆叠起来的结果 *****
16 多头注意力中，多头计算结束后的形状（堆叠）为([tgt_len,batch_size, num_heads* kdim]) torch.Size([5, 2, 32])
17 多头计算结束后，再进行线性变换时的权重W_o的形状为([num_heads*vdim, num_heads* vdim ])torch.Size([32, 32])
18 多头线性变化后的形状为([tgt_len,batch_size,embed_dim])torch.Size([5,2,32])