Transformer源码复现（学习笔记）

启发

目前Transformer应用的领域及其广泛，NLP不用提，CV，Audio Generate方面也是大放异彩，原本博主对transformer只是一知半解，大致了解算法内容，有Encoder，Decoder，还有一些attention的Q、K、V的细节，不过在不断地学习中发现，学习transformer的源码更能够去理解目前先进算法的原理，因此主动去了解了一下transformer的源码细节，这里记录一下学习过程的笔记。

基本结构

如图，transformer实际上就是图上这个玩意儿，看上去东西不多，实际实现细节还挺麻烦。

首先在打源码之前还是得看看基本架构，也就是我们需要实现什么。

通过图中最下面部分我们就能知道，inputs和outputs在进入我们的大模块前，都经过了一个embedding层，这个词嵌入和位置嵌入的知识点，可以网络搜索一下，了解一下即可。所以我们需要实现一个对我们的seq进行embedding的模块。

随后往上面看，有一个加号，那里是我们的Positional Encoding，也就是位置编码层，这是另一个模块。

再往上看，发现没有，左边的inputs对应的是N个大模块，每个大模块里面还有两个小模块。

右边的outputs对应的也是N个大模块，每个大模块里有三个小模块。

两边都看过了，我们现在看左边，为什么先看左边，我们肯定是要从输出层看起，不能吃着碗里的，看着锅里的哈。

左边我们来跟着线进去哈，发现什么，一条线散开了！四条线！，不过还好，最左边那个线我们有过深度学习基础的同学都知道，叫做short-cut，也叫residual value， 也就是残差， 这里不细讲。既然是残差，我们就要实现啊，所以这里也有一个模块。

中间三条线，我们对应着什么呢，来看Multi-Head Attention，这个是一个多头注意力机制模块，具体细节，大家看算法解读，源码解读这里不做赘述。好这个也是一个模块。

在向上，是不是有一个Add&Norm的模块，这个我们叫做Layer-Normalization， 层归一化，我们可以看出，是先norm再进行一个add，所以这个也要实现。

向上走，发现了一个Feed Forward层，这个是一个前馈神经网络，这个玩意儿主要做的是非线性变换，里面门道很多，对于初学者，知道即可。详细可以去看看算法详解。同时还包含一个shor-cut部分，这个也是一个模块。

继续往上走，发现还有一个Add&Norm的模块，与之前类似，所以就不说了。后续的short-cut部分，每一层都有，除了最后的输出层，因此不多加赘述。

好了，左边的看完了，我们来看右边的，右边的，往上面看，是不是与我们左边的第一层一样？只不过多了一个Masked，因此我们需要统一实现一个Multi-Head Attention，再在里面添加一个mask=None的参数来动态使用mask。这个模块可以和左边第一层的模块合并成一个。继续往上走，发现还有一个Add&Norm的模块，与之前类似。

在往上，我们这里有好多条线，甚至Encoder也来插了一脚，这个是什么，大家看过Transformer算法的都知道，这个是Cross-Attention，既将Encoder的特征表述层K、V放入Decoder中，然后将Decoder的Q放入，达到一个用tgt(Outputs)的Query查询Encoder中的表征，获取相似度。也就是让我们的tgt(这里简写，指代target，也就是outputs)对src(source，也就是inputs)进行查询，看看我们的最终答案跟Encoder中提取出来的特征的关系，并把得出来的关系建立在V上表达，这里不细说。继续往上走，发现还有一个Add&Norm的模块，与之前类似。

继续往上走，Feed Forward 和Add&Norm模块，都提到过了。

最后有一个Linear和SoftMax，这里学过深度学习图像部分的一些算法的，都知道我们这里再做一个分类任务，得到最终的结果，也不详细解释，因为该篇文章主要针对源码复现。

好了，说了这么多，一共有多少个模块？是不是挺多的，我们这里列下来，embedding，Positional Encoding， Multi-Head Attention，residual value，Add&Norm，Feed Forward

Cross-Attention，让我们按照顺序，一一实现吧！

分模块实现

Embedding

对于Embedding模块，torch.nn中为我们实现了该模块，因此简单调用一下即可。

对于代码，我在学习中遇到的一些注意点都加了注释。

对于Forward函数中的x， 这里提供一些思路， 我们需要预测序列是一堆seq， 这里假设batchsize是1， 我们输入到模型之前的seq类似于"i love you"， 属于一个字符串的形式。

我们需要再单词表中传入seq，将seq变成一个torch的索引， 随便举例一个[2, 10, 15]， 这是seq的索引，对于输入到模型中， 一般我们的要求是将seq转换为这样的示例

size(batch_size, seq_len, model_dim)，这里的model_dim我们是在嵌入后才生成，因此传入到嵌入层之前我们的要有(batch_size, seq_len)的size才可以进入到embedding中，随后embedding干了一个什么事情呢，

他通过生成了vocab的索引之后，变成一个(vocab_size, model_dim)的词嵌入矩阵，然后我们利用seq中的索引到里面去找对应的向量，一个索引对应一个(model_dim)的向量， 也就是一个词对应的一个向量，model_dim属于向量空间，尽可能大一点会比较好，能够捕捉到更多的语义，语法，语境特征。

找到了就将对应的索引替换成(model_dim)的向量，最后我们从embedding层得到的是一个(1, 3, model_dim)的seq，1对应的批次，3对应我们有几个单词，model_dim对应我们一个单词所能蕴含的信息量大小。

最后的缩放因子不用管，都这么写，具体看论文。

class Embeddings(nn.Module):
    def _init(self, vocab_size, d_model):  # vocab_size: 词汇表大小，d_model: 词嵌入维度
        super(Embeddings, self).__init__()  # 继承父类的初始化方法，基本上都是这么写
        self.lut = nn.Embedding(vocab_size, d_model)  
# 这里类似于生成一个embedding矩阵，size为(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)  
# 返回词嵌入矩阵乘以一个缩放因子，这个缩放因子是为了防止embedding之后的数值过大
   

Positional Encoding

上面实现了第一个模块的功能，也叫做词嵌入，但是我们得到一个词嵌入，并不能完全的知道整体的一个语境信息，例如："我做了作业"和 “作业被我做了”， 这个是两个一样的事情，但是如果不加入我的具体位置，或者作业的具体位置，模型就只能够知道我后面可以预测出作业，而作业后面也能预测出我，可能会生成“作业做了我”这种荒唐的信息。 因此位置编码(Positional Encoding)至关重要。

这里博主没有细究位置编码的细节，而是通过github上的实现源码拿出来直接使用，有兴趣的小伙伴可以探究一下位置编码的公式细节。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)  # 初始化一个(max_len, d_model)的矩阵
        position = torch.arange(0, max_len).unsqueeze(1)  # 生成一个(max_len, 1)的矩阵
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))  # 生成一个(d_model//2,)的矩阵

        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数列
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数列

        pe = pe.unsqueeze(0).transpose(0, 1)  # (max_len, d_model) -> (1, max_len, d_model)
        self.register_buffer('pe', pe)  # 将pe注册为buffer，buffer是模型的一部分，但是不会被优化器更新

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]  # x.size(0)表示x的第0维的大小，这里是一个batch的大小
        return self.dropout(x)

LayerNorm

为什么我们这里先实现LayerNorm呢，因为LayerNorm在图中出现的次数非常高频，因此我们先实现它。

在学习中我对LayerNorm的理解更深层次，LayerNorm主要是对我们输入的seq的最后一个维度，既model_dim维度来进行计算，最后通过广播机制来实现统一的计算。

我们完成了LayerNorm后就可以完成残差和LayerNorm统一的大模块了，下一个模块会讲到。现在我们先完成LayerNorm。

关于LayerNorm的公式，我放在这里，有兴趣的同学可以对着公式复现一遍代码。

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))    #features等价于model_dim
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

# a_2和b_2是可学习的参数， a_2是缩放参数，b_2是平移参数

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x- mean) / (std + self.eps) + self.b_2

SubLayerConnection

这里更改了名字，官方代码中叫SublayerConnection，主要是用于我们上文中提到的short-cut操作，以及小模块的运行操作，既一个小模块+add&norm。

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
    
    
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

学到这里博主就有个疑问，为什么要在Norm之后再来一个sublayer，明明sublayer中有norm操作。当时就查了相关资料，发现原来是在论文《Attention is All You Need》中提到过在每个子层中先利用LayerNorm再进行计算。具体可以去看论文。

实现了SublayerConnection，我们就方便的多了，因为基本上所有的模块都是基于该模块来实现的，只需要在该模块传入一个sublayer与x即可实现残差和Norm。

接下来需要来到我们的重头戏了，Self-Attention和Multi-Head Attention，建议读者先阅读相关算法详解再来阅读本文章。

Self - Attention

Self-Attention模块主要是通过注意力机制实现，具体算法在这里不仔细研究，主要研究代码如何复现。

Attention主要是通过什么呢，Q、K、V来实现，既我们给每个序列都分配好了QKV的对应向量，然后来进行公式计算，每个QKV的shape是这个样子的(batch_size, seq_len, model_dim)。

计算score的公式如下，有了这个公式，Attention模块便可以很快实现。

def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / np.sqrt(d_k)
    
    if mask is not None:
       scores = scores.masked_fill(mask == 0, 1e-9)
    p_attn = scores.softmax(dim=-1)
    
    if drop_out is not None:
        p_attn = dropout(p_attn)
    
    return torch.matmul(p_attn, value), p_attn

在这段注意力机制中，有一个mask，也就是掩码，这个可以先不看，下面会讲关于mask的操作。

对于Attention机制中的QKV是如何得来的，我才开始也不太明白，只知道如何计算，后面经过一些查询和博主的视频，发现QKV初始其实都是与weight相乘得来的，随后模型不断地学习weight，为什么要学习weight，是需要weight进行不断地优化，让模型得到不一样的特征，Q学习到的是局部的查询特征，K学习到的是全局的特征，而V学习到的是上下文的依赖特征，因为综合了QK的权值结果。

Subsequent_mask

刚刚的自注意力机制中具有一个mask，mask是什么呢？

博主在学习过程中也有这个问题，因此去查询了相关资料，发现mask其实是transformer的一个特性，即只允许现在的看见过去的，而看不见未来的，这样就能让模型具有预测未来的性能。即单向预测。与BERT不同，BERT更像是做阅读理解，向里面填词。

mask的作用我们知道了，也就是挡住一部分，不让模型看到未来的信息。同时mask还能够做一件事，也就是挡住我们的padding。

为何要挡住padding？

在词嵌入中我们说过，输入的size是(batch_size, seq_len, model_dim)形式的，那么句子有长有短，我们总要找到最长的句子，或者找到一个差不多长度的句子，将它设置为我们句子的最大长度，超出的截断，小于的padding，达到我们序列统一。而padding，我们通用的是0来padding，因此padding不含我们的信息，也就是模型不需要去学习padding所蕴含的信息，所以这里需要mask掉padding。

而position的mask原文中没有实现，应该是在输入的时候实现的，即进入到模型之前，这里实现的是subsequent_mask。

def subsequent_mask(size):    # 这里的size代表着序列的长度
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape),diagonal=1).type(torch.uint8)
    return subsequent_mask == 0
    

该段代码完成了对序列后续的mask，最后输出的mask是这样的，

是一个三角矩阵，返回的mask == 0是一个翻转过来的三角矩阵，即0与1反转。

得到这个mask后，大家再去上面看看Attention的代码就可以理解mask的含义，在attn中我们利用Q与K转置相乘得到了一个scores，scores的size是(batch_size, seq_len, seq_len)，对应着我们这边的size,通过广播机制，每一个batch都进行mask，即第一个位置只能影响到他自己，而第二个位置可以同时影响第一个和第二个，以此类推，最后一个位置可以全部都拿来加权求和。

Clone

在这里还要再来插入一个Clone模块，这个Clone模块主要实现了一个很小的功能，复制子层，N等于几就复制几个，用于多次进行模块的搭建。

def clone(modules, N):
    return nn.ModuleList([copy.deepcopy(modules) for _ in range(N)])

Multi-Head Attention

通过了这么多的模块预热，终于可以来实现目前来说最重要的一个模块，多头注意力机制了。

在学习中我才开始不知道这个多头注意力机制每个头对应的是什么，现在才发现每一个head对应的是一个seq的部分model_dim，通过N个head同时学习model_dim，能够达到更加细致的特征学习。该模块的基础是Attention，不同的是我们需要实现QKV的多维矩阵来实现该模块。

class MultiHeadAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), N=4) 
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unqueeze(1)
        nbatches = query.size(0)
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        
]
        x, self_attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        
        x= (
            x.transpose(1,2)
            .contigous()
            .view(nbatches, -1, self.h * self.d_k)
)
        del query
        del key
        del value
        self.linears[-1](x)

这个模块的实现比较难以理解首先在模块初始化，

我们需要传入一些基础参数，如h-heads，d_model-seq中每个单词最后一维的大小，dropout-去除神经元的概率。 同时初始化时，需要确保我们的d_model可以被head整除，确保每个head分得同样的features，同时得到d_k。

这里还需要做四个linear来为QKV做线性变换，最后一个linear用来储存计算出的权重结果。

在Forward函数中不好理解是一个zip的列表推导式，这里需要python基础非常扎实，zip返回出的是一个linear和对应的QKV，由于广播机制，最后一个Linear不会被提取出来。

提取出的Linear分别对QKV做线性变换，这样我们就可以通过view来更改空间维度，也就是(nbatches, seq_len, head, d_k)，这样每个头就有自己对应的d_k了，不过多头注意力机制的核心是多头！也就是采用并行计算机制，所以head和seq_len需要进行维度交换，否则无法传入到Attention模块进行计算。

这里我当时有些不理解，不过仔细想一下，如果是简单的单头自注意力机制采用的size是(batch_size, seq_len, d_model)， 如果是多头，肯定不能替换掉seq_len，因为一个头要计算整个单词序列中的一部分特征，所以head应该在seq_len之前。

拿到了对应的QKV后就可以进行Attention计算了，传入参数，得到两个返回，第一个返回是softmax后再与V相乘得到的特征，第二个是softmax后的得分值。

随后我们就可以得到新一轮的特征，经过最后一个Linear线性变换返回。

PointwiseFeedForward

前馈神经网络，对与前馈神经网络，我最开始也不太理解，在图像分类中或者语义分割中，有些模型拿它进行一个特征展平，再进行特征缩小，知道缩小到需要的类目数量，不过在这里前馈神经网络做了一个特征空间的放大，让模型能够学习融合到更多的信息，再进行特征缩小，回到原来的维度大小，这样做的好处一个是加速模型的训练和计算，另一个是可以增强模型的特征学习。

class PointwiseFeedForward(mm.Module):
    def __init__(self, d_model, d_ff, dropout=0.1)
        super(PointwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x)
        
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

现在我们基本的小模块都做完了，还剩几个大模块，比如EncoderLayer，DecoderLayer，Encoder，Decoder，Encoder_Decoder，这个几个大模块，让我们来看看如何搭建。

EncoderLayer

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout)
        super(EncoderLayer, self).__init__()
        self.attn = self_attn
        feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size
    
    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

这里博主在学习的时候也有点小问题，第一个问题就是这个size，由于表示不明确，总是容易将size搞混，不过目前出现的size，无论是在LayerNorm中的features还是这里的size，都代表着d_model，也就是每个词嵌入的向量的维度。

第二个问题就是sublayer的传入参数问题，这个可以详细搜一下lambda回调函数，在SublayerConnection中我们传入的是(x，sublayer)，sublayer就是一个子层的意思，那么在一个子层我们应该传入attn模块，来实现多头注意力机制，不过多头注意力机制需要的是四个参数，

(query, key, value, mask)， 我们直接传入attn函数的话，在SublayerConnection中也无法传入四个参数，只能传入一个norm(x)，因此这里利用lambda回调，设置参数，既将norm(x)，作为参数提前设置好，再传入当前的mask就可以了。

第二个sublayer子层就是做了一个FFN的操作。

得到了EncoderLayer，我们就要开始实现Encoder这个大模块了。

Encoder

Encoder说白了就是把一堆小模块拼在一起就可以了，其中的norm操作 我们先前也说了，在任何操作如注意力、FFN之前都需要进行一个LayerNorm的操作。

class Encoder(nn.Module):
    def __init__(self, layers, N):
        super(Encoder).__init__()
        self.layers = clone(layers, N)
        self.norm = LayerNorm(layers.size)
      
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

DecoderLayer

终于说到DecoderLayer了，关于Decoder部分，包含初始的自注意力和中间的交叉注意力，然后是最后的FFN层，所以我们就按照这个顺序来搭建一下Layer部分。

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask)
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.self_attn(x, m, m, tgt_mask))
        return self.sublayer[2](x, lambda x: self.feed_forward(x))

这里依旧利用的lambda函数。

Decoder

Decoder与Encoder类似，都是定义layers来clone N个残差归一块，再传入对应的参数来初始化。

class Decoder(nn.Module):
    def __init__(self, layer, N)
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask)
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

Generator

我们做完了Encoder，Encoder传入到Decoder，再来做Decoder，最后得到的输出还需要经过线性层，根据不同的任务，线性层的维度含义不同。

class Generator(nn.Module):
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)

EncoderDecoder

最后一步！终于可以登顶了，在搭建了这么多的小模块，中间层，大模块之后，我们终于可以把所有的模块全部都整合到一起了。

class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, src_embd, tgt_embd, generator):
        super(EncoderDecoder, x).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask):
        memory = self.encode(src, src_mask)
        res = self.decode(memory, src_mask, tgt, tgt_mask)
        return res


    def encode(self, src, src_mask):
        src_embedds = self.src_embed(src)
        return self.encoder(src_embedds, src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        target_embedds = self.tgt_embed(tgt)
        return self.encoder(target_embedds, memory, src_mask, tgt_mask)

这里我的理解就是综合了所有的EncoderDecoder功能，但是都没有进行初始化，既还未传入一些基础的参数。主要的功能是链接了EncoderDecoder。

MakeModel

我们还需要一个方法来初始化这个模型，就跟我们以前搭建CNN的分类或者分割模型一样，最后需要传入对应的参数来初始化EncoderDecoder这个类，来完成整个模型。

def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):
    c = copy.deepcopy
    attn = MultiHeadAttention(h, d_model)
    ff = PointwiseFeedForward(d_model, d_ff)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(

        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab),
)
    
    for p in model.parameters():
        if p.dim() > 1 :
            nn.init.xavier_uniform_(p)

    return model

在这里具体的Transformer源码就告一段落了，具体的代码细节可能有所疏忽，本文仅供交流学习，不做为官方参考，需要官方源码的，哈佛大学实现网址在这里：The Annotated Transformer，具体如何使用Transformer先跑起来一个样例，咱们下一篇再说。