transformer原理及代码实践

lankuohsing

已于 2023-03-11 15:55:08 修改

阅读量2.5k

点赞数 1

分类专栏：自然语言处理深度学习 pytorch 文章标签： transformer 深度学习机器翻译人工智能自然语言处理

于 2022-03-12 12:36:47 首次发布

本文链接：https://blog.csdn.net/THUChina/article/details/123441732

版权

深度学习同时被 3 个专栏收录

12 篇文章 0 订阅

订阅专栏

自然语言处理

9 篇文章 0 订阅

订阅专栏

pytorch

2 篇文章 0 订阅

订阅专栏

文章目录

0. 背景
1. transformer的宏观结构解析
2. 各模块的代码细节和数据的流动
3. 一些工具代码

参考链接和文献：
https://jalammar.github.io/illustrated-transformer/
https://blog.csdn.net/qq_28168421/article/details/120340360
http://nlp.seas.harvard.edu/annotated-transformer
完整代码及测试用例见https://github.com/lankuohsing/Study_NLP/tree/master/transformer/annotated_transformer

0. 背景

之前的基于卷积或者RNN的结构，在处理长序列数据时，计算量随着序列长度增加而增加。这样容易丢失长距离的信息。而在transformer中计算量相对于序列长度是常数。

1. transformer的宏观结构解析

transformer的一种典型的seq2seq结构，常用于序列到序列的应用中。以机器翻译为例，最顶层的视图为：
在这里插入图片描述

将结构放大一些，可以看到是input流经encoder部分再流经decoder部分最后得到output：
在这里插入图片描述

而encoder部分是由6层EncoderLayer组成，decoder部分也是由6层DecoderLayer组成，最后一层EncoderLayer会影响每一层DecoderLayer的输入：
在这里插入图片描述

每层EncoderLayer的结构都是相同的，由一层self-attention和一层FeedForward（以下简称fc）层组成：
在这里插入图片描述

其中self-attention能够使得在编码某个词时，看到并结合句子里面每个词的信息。self-attention的原理解释见https://blog.csdn.net/THUChina/article/details/108611559

DecoderLayer部分也有类似与EncoderLayer中的self-attention层和fc层，但是在这两者之间还有一个encoder-decoder attention层，用于获取原始输入句子中的相关部分。如下图所示：
在这里插入图片描述

2. 各模块的代码细节和数据的流动

2.0. 对输入的处理

对于一个序列（比如一个句子，记为x），首先要对序列中的每个元素（例如每个词）进行embedding，得到一个embedding向量：
在这里插入图片描述
在transformer里面，默认每个词的embedding向量的维度是512维。也即每个词会得到一个512维的向量表示。
具体embedding的方法是由token embedding和position embedding相加得到，前者类似于Word2Vec里面的Word Embedding，可实现预训练得到，实际使用的时候可以查表得到；后者的理论介绍见https://blog.csdn.net/THUChina/article/details/108611559 。
token embedding的代码实现为：

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(num_embeddings=vocab, embedding_dim=d_model)
        self.d_model = d_model

    def forward(self, x):
        '''
        input:
            x:一个batch的输入序列构成的矩阵，输入序列的每个元素是该元素在词表里面的index
        reurn: 上述batch的输入序列的Embedding向量构成的矩阵
        '''
        return self.lut(x) * math.sqrt(self.d_model)#归一化方差

注意，上面对得到的Embedding向量乘以了Embedding维度的根号。原因见https://www.zhihu.com/question/415263284
position embeddingg的代码实现为：

class PositionalEncoding(nn.Module):
    '''Implement the positional encoding.'''
    def __init__(self, d_model, dropout, max_len=5000):
        '''
        d_model: 输入序列中每个元素的embedding向量长度
        max_len: 用于产生positional embedding的最大长度，不能小于输入序列x的最大长度
        '''
        super(PositionalEncoding, self).__init__()#调用父类的__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)# shape: [max_len, d_model]
        position = torch.arange(0, max_len).unsqueeze(1)# value: 0~max_len; shape: from [max_len] to [max_len,1]
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )# 1/10000^(2*i/d_model)),shape[d_model/2],2*i is the even dimension
        '''
        position.shape:[max_len,1],value:[0,1,...,max_len-1]
        div_term.shape:[d_model/2],value:[10000^(2*i/d_model) for i in range(0,d_model,2)]
        shape of position * div_term:[max_len,d_model/2], value: pos/10000^(2*i/d_model)
        '''
        pe[:, 0::2] = torch.sin(position * div_term)#偶数维度
        pe[:, 1::2] = torch.cos(position * div_term)#奇数维度
        pe = pe.unsqueeze(0)# shape: [1, max_len, d_model]
        self.register_buffer("pe", pe)# not parameters. 会保存在state_dict中

    def forward(self, x):#word embedding+position embedding
        '''
        输入:
            x:输入序列的token embedding矩阵。每个元素的embedding向量长度为d_model
        返回:
            x的token embedding+positional embedding
        '''
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

2.1. Encoder部分

Encoder部分是由6层EncoderLayer堆砌而成的，每层EncoderLayer都是一个self-attention+feed-forward-neural-networks结构，Encoder的类实现如下：

"""Encoder Stacks"""
class Encoder(nn.Module):
    '''Core encoder is a stack of N EncoderLayers'''
    def __init__(self, layer, N):
        '''
        input:
            layer: EncoderLayer
            N: Num of layers. Default 6
        '''
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)# 本教程的transformer中N=6
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        '''
           Pass the input (and mask) through each layer in turn.
           输入:
               x: 一个batch输入序列x的embedding矩阵(batch_size, max_len, d_model)，这里的embedding已经是token embedding+position embedding了
               mask: 用于在attention计算的时候将padding的部分置为0
        '''
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)# 因为SublayerConnection中的code simplicity，这里要在最后加上LayerNorm操作

结合上一小节可知，Encoder的输入就是一个包含多个词的embedding向量（用x表示）的list(或者理解为一个embedding矩阵，每一行是一个元素的embedding向量)。输入可以限制一个最长长度，通常是训练集里面最长的句子长度。实际输入的序列如果小于这个最长长度，不足部分要进行padding，填充0。在代码实现上，则是通过引入mask机制(src_mask)

2.1.1. EncoderLayer

如前文所述，EncoderLayer其实就是multi-head-self-attention+feed-forward-neural-networks结构(简称sel-attenion和ffn)：
在这里插入图片描述
代码实现如下：

"""EncoderLayer"""
class EncoderLayer(nn.Module):
    '''Encoder is made up of self-attention and feed forward NN(defined below)'''
    def __init__(self, size, self_attn, feed_forward, dropout):
        '''
        size: equivalent to d_model, which is the size of embedding vector of each token
        self_atten: MultiHeadedAttention
        feed_forward: PositionwiseFeedForward
        '''
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)#这里的sublayer是self-attn或ffn（实际的计算过程在上面SublayerConnection.forward中）以及对应的残差连接和layernorm
        self.size = size

    def forward(self, x, mask):
        """
        先计算multi-head attention（以及对应的layernorm和残差连接），再计算ffn（以及对应的layernorm和残差连接）
        输入：
            x: token embedding + position embedding
            mask: 用于处理输入长度不足最大长度的padding部分
        """
        #这里的三个x分别用于后面计算query, key, value
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))#之所以用lambda表示，是因为self_atten的计算函数定义，输入是三个不同的q,k,v。但是实际它们又是一样的
        return self.sublayer[1](x, self.feed_forward)

self-attention层和FFN层各自的输入和输出通过残差连接，并且对各自输出再做一个layernorm。也即每个子结构的输出可以表示为： $L a yer N or m (x + S u b L a yer (x))$ 。其中 $x$ 代表该子层的输入， $S u b L a yer$ 代表该子层原始的实现函数（self-attention或者FFN）：
在这里插入图片描述
残差连接+layernorm的抽象代码如下：

class SublayerConnection(nn.Module):# 其实就是残差+layernorm
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    Comments: That is to say, the norm op is executed when the last output is being fed to the current layer instea of being generated from last layer. 
    So in Encoder/Decoder, there needs a LayerNorm op after the for loop of self.layers.
    """
    def __init__(self, size, dropout):
        '''
        size: equivalent to d_model, which is the size of embedding vector of each token
        '''
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):#这里是sublayer是self-attn或者fc，对应下文的lambda表达式
        '''
        Apply residual connection to any sublayer with the same size.
        理论上的计算公式为：LayerNorm(x+SubLayer(x))。但是在代码实现上为x+Dropout(SubLayer(LayerNorm(x)))。
        先对上一层的输出x做layernorm，然后进行sublayer计算（attention或者Feedforward），然后dropout，然后残差连接。
        所以本函数在最后是少算了一个layernorm的。
        开头对原始的输入embedding也算了一次layernorm，跟图里面不一致
        x: shape: (batch_size, max_len, d_model)
        '''
        return x + self.dropout(sublayer(self.norm(x)))#真正的self-attn或者fc的计算发生在这里

。而self-attention的输出也是一个list，里面包含了每个词经过该层self-attention后的representation vector（用z表示），attention具体计算过程参见https://blog.csdn.net/THUChina/article/details/108611559 ：

然后每个词的z向量，又经过一层全连接层。
此外，分别对每层EncoderLayer中的
同时，为了方便残差连接的实现，每层EncoderLayer的输入输出维度保持一致（512维），这样就不需要在残差连接里面对 $x$ 进行升维或者降维。
注意，在Encoder部分里，某层attention计算过程可以关注上一层输出的所有位置的信息。
Encode里面的mask机制(对应位置的softmax输入替换 $-\infty$ )，是为了掩盖序列中的padding部分的值的，防止它们参与计算。

2.2. Decoder部分

注意，在Encoder部分中，无论输入序列中元素个数有多少，前向过程都可以并行地对输入序列的embedding矩阵机械能矩阵乘法运算，这也是self-attention相比RNN/LSTM的优势，可以并行加速；而Decoder部分中，必须进行串行地循环推理，逐个生成新的输出元素（单词）。当然，在训练过程中，Decoder可以通过teacher force和masked self attention来实现并行
decoder部分的结构图如下，可以看到比encoder多了一个attention：
在这里插入图片描述
代码如下：

class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        '''
        input:
            layer: DecoderLayer
            N: Num of layers. Default 6
        '''
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        '''
           Pass the input (and mask) through each layer in turn.
           输入:
               x: decoder目前已经产生的序列对应的embedding矩阵，(batch_size, seq_len,d_model)
            memory：encoder的输出embedding矩阵，，(batch_size, seq_len,d_model)
            src_mask: encoder的mask,用于在attention计算的时候将padding的部分置为0
            tag_mask:为了防止decoder看到未来信息的mask，是一个上三角矩阵
        '''
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

2.2.1. DecoderLayer

注意，DecoderLayer的第一个self-attention的计算过程，与EncoderLayer中的self-attention的计算过程类似（query，key，value的计算都是依赖于上一层的输出，具体计算过程参见https://blog.csdn.net/THUChina/article/details/108611559 ），而第二个self-attention中的key和value来自Encoder的输出计算得到，query来自上一个子层的输出计算得到。可以这么形象地理解，第二个self-attention层的功能是利用解码器已经预测出的信息作为query，去编码器提取的各种特征中，查找相关信息并融合到当前特征中，来完成预测。代码如下：

"""DecoderLayer"""
class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn# (q,k,v,mask,dropot)
        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):
        """
        先计算masked multi-head self-attention（以及对应的layernorm和残差连接），再计算multi-head encoder-decoder attention（以及对应的layernorm和残差连接）, 再计算ffn（以及对应的layernorm和残差连接）
        输入：
            x: decoder目前已经产生的序列对应的embedding矩阵，(batch_size, seq_len,d_model)
            memory：encoder的输出embedding矩阵，，(batch_size, seq_len,d_model)
            src_mask: encoder的mask,用于在attention计算的时候将padding的部分置为0
            tag_mask:为了防止decoder看到未来信息的mask，是一个上三角矩阵
        """
        m = memory
        # self-attention的q,k,v都是通过上一层输入计算得到
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        # encoder-decoder-atention的q是通过上一层输入计算得到，k,v是通过encode输出计算得到
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

注意，第一个self-attention中，在计算当前位置p的attention值时，只能关注历史信息，需要对上一层输出序列里面超过当前位置的值进行mask(对应位置的softmax输入替换 $-\infty$ )
第二个self-attention不需要进行mask，也即计算每个位置的attention值时可以关注Encoder输出中的所有位置信息。

2.3. EncoderDecoder整体结构

"""Model Architecture: Encoder-Decoder"""
class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        '''
        src_embed:包含了token embedding和position embedding。也即src_embed已经是前面两个embedding的sum了
        '''
        super(EncoderDecoder, self).__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):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask,
                            tgt, tgt_mask)

    def encode(self, src, src_mask):
        '''
        encoder部分的计算过程
        src: 输入序列（包含序列中各个元素的下标）.shape: (batch_size, max_len)
        src_mask: 用于掩盖padding部分的mask,ByteTensor类型。shape:(batch_size,1,max_len)
        '''
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        '''
        x: decoder目前已经产生的序列对应的embedding矩阵，(batch_size, seq_len,d_model)
        memory：encoder的最后一层输出embedding矩阵，，(batch_size, seq_len,d_model)
        src_mask: encoder中的mask
        tgt：decoder的输出序列的元素序号？初始为[0], shape:(1,seq_len+1)
        tgt_mask:上三角矩阵，右上角为false,左下角（包括对角线）为true.shape:(1,seq_len+1)
        '''
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

3. 一些工具代码

3.1. 将一个模块复制N份，注意要深拷贝:

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

3.2. layernorm计算：

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        '''
        features: d_model, size of embedding vector
        '''
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        '''
        x: 上一个sublayer的输出，最后一个维度为序列中元素的特征维度
        '''
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

3.3. attention机制相关代码

3.3.1. self-attention

def attention(query, key, value, mask=None, dropout=None):
    """
    Compute 'Scaled Dot Product Attention'
    if multi-head,then shape of q/k/v is (nbatches, self.h,max_len, self.d_k)
    d_k=d_model // h
    return: attn_outputs.shape: (nbatches, self.h, max_len, self.d_k)
            p_attn.shape: (nbatches, self.h, max_len, max_len)
    """
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)# scores.shape (nbatches, self.h,max_len, max_len)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)# mask为0的索引，在scores中同样索引处置为负无穷，这样经过softmax后权重为0
    p_attn = scores.softmax(dim=-1)# p_attn.shape (nbatches, self.h, max_len, max_len)
    '''对序列中每个元素i，它与序列中每个元素的相关性分数，所以一个序列的p_attn是一个(max_len, max_len)的矩阵'''
    if dropout is not None:
        p_attn = dropout(p_attn)
    attn_outputs=torch.matmul(p_attn, value)# attn_outputs.shape (nbatches, self.h, max_len, self.d_k)
    '''计算序列中每个元素的representation vector，需要用该元素的attention分数乘以序列中每个元素的value vector，所以得到的是一个(max_len, self.d_k)的矩阵'''
    return attn_outputs, p_attn

3.3.2. multi-attention

"""multi-head attention"""
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):# d_model=512, h=8
        '''
        h: Number of heads. Default 8
        d_model: model size (size of embedding vector)
        '''
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h# 64
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        '''
        这里的q,k,v其实都是上一层的输出向量，所以是一样的，真正的q,k,v由下面计算得到
        shape of x: [batch_size, max_len, d_model], max_len是序列最大长度，默认为512（序列长度不足部分padding 0，所以其实各个样本的len_seq是一样的，都是max_len）
        '''
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        # 1) Do all the linear projections in batch from d_model => h x d_k
        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))
        ]#分别对q,k,v乘以权重矩阵。只用到了3个linear layer.这是是有bias的
        # q,k,v的shape: (batch, max_len, d_model)->(batch,h,max_len,d_k)

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )# x.shape: (nbatches, self.h, max_len, self.d_k) self.attn.shape: (nbatches, self.h, max_len, max_len)

        # 3) "Concat" using a view and apply a final linear.
        x = (
            x.transpose(1, 2)
            .contiguous()# deepcopy
            .view(nbatches, -1, self.h * self.d_k)
        )# x.shape: (nbatches, max_len, d_model)
        del query
        del key
        del value
        multi_head_output=self.linears[-1](x)#用上最后一个linear layer
        # multi_head_output.shape: (batch,max_len,d_model)
        return multi_head_output

3.4. 全连接网络

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation. fc->relu->fc"

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, 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(self.w_1(x).relu()))

3.5. decoder中的mask

为了防止看到未来信息

def subsequent_mask(size):
    "Mask out subsequent positions. 用于DecoderLayer中的self-attention sub-layer，以防止看到未来的信息"
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )#保留torch.ones(attn_shape)的上三角矩阵，diagonal=1表示不保留对角线
    return subsequent_mask == 0

3.6. 对decoder的输出embedding计算概率

class Generator(nn.Module):
    "Define standard linear + softmax generation step."
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(in_features=d_model, out_features=vocab, bias=True)

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

lankuohsing

关注

1
点赞
踩
11

收藏

觉得还不错? 一键收藏
1
评论
transformer原理及代码实践

文章目录1. transformer的宏观结构解析2. 数据的流动2.1. Encoder部分2.2. Decoder部分3. 代码实现（待补充）参考链接和文献：https://jalammar.github.io/illustrated-transformer/https://blog.csdn.net/qq_28168421/article/details/1203403601. transformer的宏观结构解析transformer的一种典型的seq2seq结构，常用于序列到序列的应用.
复制链接

扫一扫