Attention Is All You Need，Transformer代码实现

这篇论文很早就读过，当时只是简单了解了下其原理，但真正动手实现时还是能发现不少不能忽略的细节问题，这里不说明原理（原理都在文献上），只注重实现。

目录

架构

Transfomer

encoder

 注意力层

前馈神经网络

decoder

---

架构

 上图便是Transformer的架构，可以看到，它可以分为三个部分，encoder，decoder和Linear projection。encoder又可以分为三个部分：词向量编码，位置编码，多头注意力层和前馈神经网络层（这里做一个部分）。decoder相比encoder多了一个Masked的注意力层，其余和encoder的结构一样。如下图所示

补充：该模型有3个输入，分别是Inputs，Outputs shifted right和Output Probabilities。Inputs作为编码器的输入，Outputs shifted right作为解码器的输入，Output Probabilities用来与解码器的输出做损失，一般Outputs shifted right=Output Probabilities。

Transfomer

先设置好transformer的参数

    #参数
    Embedding_size = 512 #词向量维度
    FF_dimension = 2048 #前馈层维度
    K_size = Q_size = V_size = 64 #注意力的QKV的维度
    n_layers = 6 #编码器与解码器的层数
    n_head = 8 #多头注意力的头数

整体架构清晰了，接下来就是一个框架一个框架去实现，首先把大框架Transformer 列出来

class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        self.Linear_projection = nn.Linear(Embedding_size, length_tgv)

 因为Transformer最后的输出要与Output Probabilities做损失，那么需要保证维度一致，所以Linear_projection层要将Embeeding_size（词向量维度）的向量转为length_tgv（output词表的维度）维度。接下来实现encoder与decoder。

encoder

根据架构，很容易初始化

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.embed = nn.Embedding(length_srcv, Embedding_size)
        self.position = PositionalEncoding(Embedding_size)
        self.layers = nn.ModuleList(EncoderLayer() for _ in range(n_layers))

 这里用了EncoderLayer将注意力和前馈神经网络合在一起（当然还有Add&Norm操作）

class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.Multi_attn = MultiHeadAttention()
        self.FeedForward = FeedForwardPosition()

 注意力层

class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(Embedding_size, Q_size*n_head)
        self.W_K = nn.Linear(Embedding_size, K_size*n_head)
        self.W_V = nn.Linear(Embedding_size, V_size*n_head)
        self.Linear = nn.Linear(Q_size*n_head, Embedding_size)
        self.norm = nn.LayerNorm(Embedding_size)

 这里的W_Q，W_K，W_V，是根据文献的多头注意力公式得来的

原始的Q，K，V需要经过W_Q，W_K，W_V线性层，然后再通过注意力机制，因为等会要分头分别计算注意力，所以通过这些线性层他们的维度会变成Q_size*n_head（方便等会分成[Q_size，n_head]这样的维度。

    def forward(self, Q, K, V, attn_mask):
        temp = Q
        batch_size, len_seq, _ = Q.size()
        s_q = self.W_Q(Q).unsqueeze(1).reshape(batch_size, n_head, len_seq, Q_size) #s_q:[1,8,6,64]
        s_k = self.W_K(K).unsqueeze(1).reshape(batch_size, n_head, len_seq, K_size) #s_k:[1,8,6,64]
        s_v = self.W_V(V).unsqueeze(1).reshape(batch_size, n_head, len_seq, V_size) #s_v:[1,8,6,64]

        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_head, 1, 1) #attn_mask:[1,8,6,6]

        attn_score = ScaledDotProductAttention()(s_q,s_k,s_v,attn_mask) #attn_score:[1,8,6,64]
        attn_score = attn_score.transpose(1,2).squeeze(2).reshape(batch_size,len_seq,Q_size*n_head) #attn_score:[1,6,512]
        attn_score = self.Linear(attn_score) #attn_score:[1,6,512]
        attn_score = self.norm(attn_score+temp) #attn_score:[1,6,512]
        return attn_score #attn_score:[1,6,512]

这里注意attn_mask，这是一个矩阵，用来标识输入的句子的填充情况（pad），比如下面这样的例子：

 如果限制最大的输入句子的长度是7，那么很明显第一个和第三个句子不够长，那么就要进行填充pad，但是这些填充字符是没有意义的，注意力机制不需要进行计算，所以这个矩阵记录了哪些是句子，哪些是填充字符，比如设置如果字是1，填充是0，那么attn_mask矩阵应该是这样的：

得到s_q，s_k，s_v之后， 便根据论文给出的计算公式进行注意力点积计算

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()
    def forward(self,s_q,s_k,s_v,attn_mask):
        score = torch.matmul(s_q,s_k.transpose(2,3)) / np.sqrt(K_size)
        score = score.masked_fill_(attn_mask, -1e9)
        score = nn.Softmax(dim=1)(score)
        score = torch.matmul(score,s_v)
        return score #score[1,8,6,64]

这里使用masked_fill_函数对score根据attn_mask的形状进行赋值，如果是填充的话会将其对应的位置赋值为-1e9，因为数值太小，在计算过程中不会造成影响。

前馈神经网络

只使用两个卷积层，注意输出的结果需要与原输入相加（见结构图），最后进行归一化

class FeedForwardPosition(nn.Module):
    def __init__(self):
        super(FeedForwardPosition, self).__init__()
        self.conv1 = nn.Conv1d(Embedding_size,FF_dimension, kernel_size = 1)
        self.conv2 = nn.Conv1d(FF_dimension,Embedding_size, kernel_size = 1)
        self.norm = nn.LayerNorm(Embedding_size)
    def forward(self,attn_score):
        temp = attn_score
        output= self.conv1(attn_score.transpose(1,2))
        output = self.conv2(output).transpose(1,2)
        output = self.norm(output + temp)
        return output #output:[1,6,512]

decoder

根据结构，先将decoder初始化

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.embed = nn.Embedding(length_tgv, Embedding_size)
        self.position = PositionalEncoding(Embedding_size)
        self.layers = nn.ModuleList(DecoderLayer() for _ in range(n_layers))

这里使用DecoderLayer将注意力和前馈网络合在一起，这里要注意的是，decoder比encoder多了一个mask注意力层

class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.mask_Mutil_attn = MultiHeadAttention()
        self.Mutil_attn = MultiHeadAttention()
        self.FeedForward = FeedForwardPosition()

这个mask注意力层的作用是：它需要对经过mask的decoder输入进行注意力计算。

为什么要mask decoder的输入呢？以机器翻译为例，翻译“我要回家了”，那么encoder的input便是

“我要回家了”，decoder的输入是“i have to go home”，如果不mask的话，attention会从上下文得到答案，输入i的时候，它很容易通过下文得到have这个正确答案（有点作弊的意思），所以每次输入的时候，都要对后面的单词mask，这样就形成了一个上三角矩阵，设1为原单词，0为mask，则

 将这个上三角矩阵与填充矩阵attn_mask相加，便得到一个既有mask也有识别pad的的矩阵，再用其计算注意力。Decoder的其余部分与Encoder一样，重复调用即可。

encoder和decoder都实现了，将他们在Transformer中整合到一起，最后实现Transformer

    def forward(self, Encoder_input, Decoder_input):
        Encoder_output = self.encoder(Encoder_input)
        decoder_output = self.decoder(Encoder_output,Decoder_input)
        output = self.Linear_projection(decoder_output)

完整代码：https://github.com/bowspider-man/transformer