devbox
我家小花儿
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【深度学习】(2) Transformer 网络解析，代码复现，附Pytorch完整代码_transformer代码复现

作者：我家小花儿 | 2024-03-31 18:52:37

赞

踩

transformer代码复现

大家好，今天和各位分享一下如何使用 Pytorch 构建 Transformer 模型。

本文的重点在代码复现，部分知识点介绍的不多，我会在之后的四篇博文中详细介绍 Encoder，Decoder，(Mask)MutiHeadAttention，以及实战案例。

之前我也介绍过 Vision Tranformer，该模型的 Pytorch 和 TensorFlow2 的复现和实战案例可以看我的这两篇博文：

https://blog.csdn.net/dgvv4/article/details/124792386

https://blog.csdn.net/dgvv4/article/details/125184340

1. 引言

在传统的记忆网络中，RNN 模型，LSTM 模型以及 GRU 模型三个模型都是时序模型，对于每一个句子只能串行运行而不能并行，限制了模型的速度，而且虽然在一定程度上解决了长期依赖问题，但是对于跨度特别长的依赖还是没有很好的解决。

Transformer 模型内的每一个单词都可以并行的与其他单词计算Self-Attention值，并且在计算时不会受到距离过远的影响。Transformer 模型是一个序列到序列的模型，它由两个部分组成，分别是编码器和解码器。Transformer 也是一种基于注意力机制的架构，能够利用对文本的注意力权重来提取重要讯息。Transformer 具有可以并行处理顺序数据的特点，所以，它不但速度比以前的体系结构快，而且，在处理长期依赖性方面也十分优秀。

2. Muti_head_attention

自注意力层的主要功能是收集到有关句子的上下文的语意信息。

对于一个输入，自注意力机制首先会将 Embedding 向量与随机矩阵计算乘积得到三个向量矩阵，分别是 查询 Q（Query）、键 K（key）、值 V（value）权重矩阵。之后根据这三个向量矩阵计算自注意力的分数值。

计算过程如下：首先，将 Q 与 K 做点乘；之后将点乘的结果除以一个常数 ​，目的是为了保证计算的稳定性；然后把得到的结果经过 Softmax 层进行计算；最后，将 Softmax 的输出结果与 V 矩阵进行相乘，得到该点的注意力分数。具体公式如下：

attention(Q,K,V)=sofymax(\frac{QK^{T}}{\sqrt{d_k} } )V

其中 Q，K 和 V 是由输入词向量构成的矩阵，​dk 是输入词向量的维数。计算过程如下图所示：

从实质上可以了解，注意力机制从大量的数据中有过滤出了一些可用的关键信息，因而集中于这部分关键信息上，而忽略了这些无用的不重要的内容。softmax 计算后的权重越大，注意力就越集中在其相应的 Value 值上，即权重就代表了数据的意义，而 Value 值也就是与它意义相对应的数据。 

Transformer 模型所采用是并不是单一的自注意力机制，而是多头注意力机制。通俗的解释就是说对于一个输入向量，不仅仅只初始化一组 Q、K、V 的矩阵，而是初始化多组，最终得到多组向量矩阵，不同的头可以表示不同的语义信息，示意图如下：

Transformer 模型中采用了 8 组权重矩阵，也就是 8 个注意力头。为了满足前馈层输入要求，定义一个权重矩阵 W0 与 8 组权重矩阵相乘，对模型进行联合训练。最终将融合所有注意力头信息的矩阵送入前馈层进行下一步的计算。

多头注意力计算公式如下，其中 W_i^Q、W_i^K、W_i^V 是 Q、K、V 矩阵的参数矩阵，W^0 是附加矩阵。

head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)

MutiHead(Q,K,V)=Concat(head_1, ..., head_n)W^0

代码如下：

  1. import torch
  2. from torch import nn
  3.  
  4. # -------------------------------------------------- #
  5. #（1）muti_head_attention
  6. # -------------------------------------------------- #
  7. '''
  8. embed_size: 每个单词用多少长度的向量来表示
  9. heads: 多头注意力的heads个数
  10. '''
  11. class selfattention(nn.Module):
  12.     def __init__(self, embed_size, heads):
  13.         super(selfattention, self).__init__()
  14.  
  15.         self.embed_size = embed_size
  16.         self.heads = heads
  17.         # 每个head的处理的特征个数
  18.         self.head_dim = embed_size // heads
  19.  
  20.         # 如果不能整除就报错
  21.         assert (self.head_dim * self.heads == self.embed_size), 'embed_size should be divided by heads'
  22.  
  23.         # 三个全连接分别计算qkv
  24.         self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
  25.         self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
  26.         self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
  27.  
  28.         # 输出层
  29.         self.fc_out = nn.Linear(self.head_dim * self.heads, embed_size)
  30.  
  31.     # 前向传播 qkv.shape==[b,seq_len,embed_size]
  32.     def forward(self, values, keys, query, mask):
  33.         
  34.         N = query.shape[0]  # batch
  35.         # 获取每个句子有多少个单词
  36.         value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]
  37.  
  38.         # 维度调整 [b,seq_len,embed_size] ==> [b,seq_len,heads,head_dim]
  39.         values = values.reshape(N, value_len, self.heads, self.head_dim)
  40.         keys = keys.reshape(N, key_len, self.heads, self.head_dim)
  41.         queries = query.reshape(N, query_len, self.heads, self.head_dim)
  42.  
  43.         # 对原始输入数据计算q、k、v
  44.         values = self.values(values)
  45.         keys = self.keys(keys)
  46.         queries = self.queries(queries)
  47.  
  48.         # 爱因斯坦简记法，用于张量矩阵运算，q和k的转置矩阵相乘
  49.         # queries.shape = [N, query_len, self.heads, self.head_dim]
  50.         # keys.shape = [N, keys_len, self.heads, self.head_dim]
  51.         # energy.shape = [N, heads, query_len, keys_len]
  52.         energy = torch.einsum('nqhd, nkhd -> nhqk', [queries, keys])
  53.  
  54.         # 是否使用mask遮挡t时刻以后的所有q、k
  55.         if mask is not None:
  56.             # 将mask中所有为0的位置的元素，在energy中对应位置都置为 －1*10^10
  57.             energy = energy.masked_fill(mask==0, torch.tensor(-1e10))
  58.  
  59.         # 根据公式计算attention, 在最后一个维度上计算softmax
  60.         attention = torch.softmax(energy/(self.embed_size**(1/2)), dim=3)
  61.         
  62.         # 爱因斯坦简记法矩阵元素，其中query_len == keys_len == value_len
  63.         # attention.shape = [N, heads, query_len, keys_len]
  64.         # values.shape = [N, value_len, heads, head_dim]
  65.         # out.shape = [N, query_len, heads, head_dim]
  66.         out = torch.einsum('nhql, nlhd -> nqhd', [attention, values])
  67.         
  68.         # 维度调整 [N, query_len, heads, head_dim] ==> [N, query_len, heads*head_dim]
  69.         out = out.reshape(N, query_len, self.heads*self.head_dim)
  70.  
  71.         # 全连接，shape不变
  72.         out = self.fc_out(out)
  73.         return out

3. Muti_head_attention + FFN

前馈神经网络（Feed Forward Network，FFN）的作用主要是提供非线性变换，将多头注意力层的输出结果映射到更高维度的空间中。公式如下，可以看出就是先经过一个线性变换，再经过一个RELU激活函数，然后再经过一个线性变换。

​FFN=max(0,xW_1, b_1)W_2 + B_2

在 Transformer 模型中，每一个子层之后都会有一个残差模块，并且有一个层归一化计算（Layer Normalization，LN）。残差模块的作用主要是为了解决随着网络深度的增加梯度不稳定的问题。归一化的作用主要是防止梯度消失或者梯度爆炸，加速模型收敛。Transformer 模型中采用的是层归一化方法，计算公式如下：

​LN(x_i)= \alpha * \frac{x_i-\mu_L}{\sqrt{\sigma_L^2+\varepsilon }}+ \beta

该部分代码对应的结构图如下：

  1. # -------------------------------------------------- #
  2. #（2）multi_head_attention + FFN
  3. # -------------------------------------------------- #
  4. '''
  5. embed_size: wordembedding之后, 每个单词用多少长度的向量来表示
  6. heads: 多头注意力的heas个数
  7. drop: 杀死神经元的概率
  8. forward_expansion:  在FFN中第一个全连接上升特征数的倍数
  9. '''
  10. class TransformerBlock(nn.Module):
  11.     def __init__(self, embed_size, heads, dropout, forward_expansion):
  12.         super(TransformerBlock, self).__init__()
  13.  
  14.         # 实例化自注意力模块
  15.         self.attention = selfattention(embed_size, heads)
  16.  
  17.         # muti_head之后的layernorm
  18.         self.norm1 = nn.LayerNorm(embed_size)
  19.         # FFN之后的layernorm
  20.         self.norm2 = nn.LayerNorm(embed_size)
  21.  
  22.         # 构建FFN前馈型神经网络
  23.         self.feed_forward = nn.Sequential(
  24.             # 第一个全连接层上升特征个数
  25.             nn.Linear(embed_size, embed_size * forward_expansion),
  26.             # relu激活
  27.             nn.ReLU(),
  28.             # 第二个全连接下降特征个数
  29.             nn.Linear(embed_size * forward_expansion, embed_size)
  30.         )
  31.  
  32.         # dropout层随机杀死神经元
  33.         self.dropout = nn.Dropout(dropout)
  34.     
  35.     # 前向传播, qkv.shape==[b,seq_len,embed_size]
  36.     def forward(self, value, key, query, mask):
  37.  
  38.         # 计算muti_head_attention
  39.         attention = self.attention(value, key, query, mask)
  40.         # 输入和输出做残差连接
  41.         x = query + attention
  42.         # layernorm标准化
  43.         x = self.norm1(x)
  44.         # dropout
  45.         x = self.dropout(x)
  46.  
  47.         # FFN
  48.         ffn = self.feed_forward(x)
  49.         # 残差连接输入和输出
  50.         forward = ffn + x
  51.         # layernorm + dropout
  52.         out = self.dropout(self.norm2(forward))
  53.  
  54.         return out

4. Encoder

该部分代码对应的网络结构如下图所示，Encoder 部分堆叠了多个 TransformerBlock（Muti_head_attention 和 FFN）。输入的序列先经过 WordEmbedding，将句子中的每个单词用长度为 embed_size 的向量来表示。

由于注意力机制更多的是关注词与词之间的重要程度，而不关心句子中词语位置的顺序关系。

例如：从北京开往济南的列车”与“从济南开往北京的列车”，词向量表示并不能对两句话中的“北京”进行区分，其编码是一样的。但是在真实语境中，两个词语所表达的语义并不相同，第一个表示的是起始站，另一个表示的是终点站，两个词所表达的语义信息并不相同。

Transformer 模型通过对输入向量额外添加位置编码PositionEmbedding来解决这个问题。Transformer 模型中利用正弦和余弦函数来生成位置编码信息，将位置编码信息与 WordEmbedding 的结果相加，作为输入送到下一层。

但是下面的代码中我没用正余弦位置编码，位置编码部分我在下一篇博文中详细讲。

  1. # -------------------------------------------------- #
  2. #（3）encoder
  3. # -------------------------------------------------- #
  4. '''
  5. src_vocab_size: 一共有多少个单词
  6. num_layers: 堆叠多少层TransformerBlock
  7. device: GPU or CPU
  8. max_len: 最长的一个句子有多少个单词
  9. embed_size: wordembedding之后, 每个单词用多少长度的向量来表示
  10. heads: 多头注意力的heas个数
  11. drop: 在muti_head_atten和FFN之后的dropout层杀死神经元的概率
  12. forward_expansion:  在FFN中第一个全连接上升特征数的倍数
  13. '''
  14. class Encoder(nn.Module):
  15.     def __init__(self, src_vocab_size, num_layers, device, max_len,
  16.                  embed_size, heads, dropout, forward_expansion):
  17.         super(Encoder, self).__init__()
  18.  
  19.         self.embed_size = embed_size
  20.         self.device = device
  21.  
  22.         # wordembedding 将每个单词用长度为多少的向量来表示
  23.         self.word_embedding = nn.Embedding(src_vocab_size, embed_size)
  24.         # 对每一个单词的位置编码
  25.         self.position_embedding = nn.Embedding(max_len, embed_size)
  26.  
  27.         self.dropout = nn.Dropout(dropout)
  28.  
  29.         # 将多个TransformerBlock保存在列表中
  30.         self.layers = nn.ModuleList(
  31.             [TransformerBlock(embed_size, heads, dropout, forward_expansion)
  32.             for _ in range(num_layers)]
  33.         )
  34.  
  35.         
  36.     # 前向传播x.shape=[batch, seq_len]
  37.     def forward(self, x, mask):
  38.         # 获取输入句子的shape
  39.         N, seq_len = x.shape
  40.  
  41.         # 为每个单词构造位置信息, 并搬运到GPU上
  42.         position = torch.arange(0, seq_len).expand(N, seq_len).to(self.device)
  43.  
  44.         # 将输入的句子经过wordembedding和位置编码后相加 [batch, seq_len, embed_size]
  45.         out = self.word_embedding(x) + self.position_embedding(position)
  46.         # dropout层
  47.         out = self.dropout(out)
  48.  
  49.         # 堆叠多个TransformerBlock层
  50.         for layer in self.layers:
  51.             out = layer(out, out, out, mask)
  52.  
  53.         return out

5. DecoderBlock

该部分代码对应的网络结构如下，大体结构和 Encoder 部分相同。两点不同，首先第一个的 Muti_head_attention 中使用了掩码 mask；其次第二个 Muti_head_attention 中的 query 使用的是目标序列，key 和 value 使用的是编码器的输出序列。这部分涉及的内容较多，具体原理留到下一篇文章。

代码如下：

  1. # -------------------------------------------------- #
  2. #（4）decoder_block
  3. # -------------------------------------------------- #
  4. '''
  5. embed_size: wordembedding之后, 每个单词用多少长度的向量来表示
  6. heads: 多头注意力的heas个数
  7. drop: 在muti_head_atten和FFN之后的dropout层杀死神经元的概率
  8. forward_expansion:  在FFN中第一个全连接上升特征数的倍数
  9. '''
  10. class DecoderBlock(nn.Module):
  11.     def __init__(self, embed_size, heads, forward_expansion, dropout):
  12.         super(DecoderBlock, self).__init__()
  13.  
  14.         # 实例化muti_head_attention
  15.         self.attention = selfattention(embed_size, heads)
  16.         # 实例化TransformerBlock
  17.         self.transformer_block = TransformerBlock(embed_size, heads, dropout, forward_expansion)
  18.  
  19.         # 第一个muti_head_atten之后的LN和Dropout
  20.         self.norm = nn.LayerNorm(embed_size)
  21.         self.dropout = nn.Dropout(dropout)
  22.  
  23.     # 前向传播
  24.     def forward(self, x, value, key, src_mask, trg_mask):
  25.         
  26.         # 对output计算self_attention
  27.         attention = self.attention(x, x, x, trg_mask)
  28.         # 残差连接
  29.         query = self.dropout(self.norm(attention + x))
  30.  
  31.         # 将encoder部分的k、v和decoder部分的q做TransformerBlock
  32.         out = self.transformer_block(value, key, query, src_mask)
  33.         return out

6. Decoder

该部分代码对应网络结构如下图所示，这部分除了上面说的两点之外，其他都和 Encoder 部分相同。首先通过 WordEmbedding 将目标序列的每个单词用长度为 embed_size 的向量来表示。经过多个 DecoderBlock 之后，输入和输出的shape保持不变 [batch, seq_len, embed_size]，最后在输出层使用一个全连接层得到结果。

代码如下：

  1. # -------------------------------------------------- #
  2. #（5）decoder
  3. # -------------------------------------------------- #
  4. '''
  5. trg_vocab_size: 目标句子的长度
  6. num_layers: 堆叠多少个decoder_block
  7. max_len: 目标句子中最长的句子有几个单词
  8. device: GPU or CPU
  9. embed_size: wordembedding之后, 每个单词用多少长度的向量来表示
  10. heads: 多头注意力的heas个数
  11. drop: 在muti_head_atten和FFN之后的dropout层杀死神经元的概率
  12. forward_expansion:  在FFN中第一个全连接上升特征数的倍数
  13. '''
  14. class Decoder(nn.Module):
  15.     def __init__(self, trg_vocab_size, num_layers, device, max_len,
  16.                  embed_size, heads, forward_expansion, dropout):
  17.         super(Decoder, self).__init__()
  18.  
  19.         self.device = device
  20.  
  21.         # trg_vocab_size代表目标句子的单词总数，embed_size代表每个单词用多长的向量来表示
  22.         self.word_embedding = nn.Embedding(trg_vocab_size, embed_size)
  23.         # 位置编码，max_len代表目标句子中最长有几个单词
  24.         self.position_embeddimg = nn.Embedding(max_len, embed_size)
  25.  
  26.         # 堆叠多个decoder_block
  27.         self.layers = nn.ModuleList(
  28.             [DecoderBlock(embed_size, heads, forward_expansion, dropout)
  29.              for _ in range(num_layers)]
  30.         )
  31.  
  32.         # 输出层
  33.         self.fc_out = nn.Linear(embed_size, trg_vocab_size)
  34.         self.dropout = nn.Dropout(dropout)
  35.     
  36.     # 前向传播
  37.     def forward(self, x, enc_out, src_mask, trg_mask):
  38.         
  39.         # 获取decoder部分输入的shape=[batch, seq_len]
  40.         N, seq_len = x.shape
  41.  
  42.         # 位置编码
  43.         positions = torch.arange(0, seq_len).expand(N, seq_len).to(self.device)
  44.  
  45.         # word_embedding和位置编码后的结果相加
  46.         x = self.word_embedding(x) + self.position_embeddimg(x)
  47.         x = self.dropout(x)
  48.  
  49.         # 堆叠多个DecoderBlock, 其中它的key和value是用的encoder的输出 [batch, seq_len, embed_size]
  50.         for layer in self.layers:
  51.             x = layer(x, enc_out, enc_out, src_mask, trg_mask)
  52.  
  53.         # 输出层
  54.         out = self.fc_out(x)
  55.         return out

7. 模型构建

完整的 Transformer 模型结构图如下，与大多数端到端的模型一样，Transformer 结构也是采用了端到端的编码器-解码器（Encoder-Decoder）架构。模型将注意力的思想发挥到了极致，编码层由 6 个编码器堆叠而成，解码层由 6 个解码器组成。结构保持一致。

对于每一个编码器，主要包含两层结构。一个是自注意力层（self_attention），自注意力层的作用是获取到句子的上下文语义信息；另一个是前馈神经网络（FFN）。此外，在每个层与层之间都会有一个残差模块和归一化操作，其作用是加速模型收敛、防止梯度消失或梯度爆炸。解码器也包含编码器中的两层网络结构，除此之外，为了获取当前节点的关注信息，在这两层结构中间还有一层编码器-解码器注意力层。 

下面只需要将上面定义的 Encoder 类和 Decoder 类拼接在一起就行了。

  1. # -------------------------------------------------- #
  2. #（6）模型构建
  3. # -------------------------------------------------- #
  4. class Transformer(nn.Module):
  5.     def __init__(self, 
  6.                  src_vocab_size,
  7.                  trg_vocab_size,
  8.                  src_pad_idx,
  9.                  trg_pad_idx,
  10.                  embed_size=512,
  11.                  num_layers=6,
  12.                  forward_expansion=4,
  13.                  heads=8,
  14.                  dropout=0,
  15.                  device='cuda',
  16.                  max_len=100):
  17.         super(Transformer, self).__init__()
  18.  
  19.         self.encoder = Encoder(src_vocab_size, 
  20.                                num_layers, 
  21.                                device, 
  22.                                max_len,
  23.                                embed_size, 
  24.                                heads, 
  25.                                dropout, 
  26.                                forward_expansion)
  27.  
  28.         self.decoder = Decoder(trg_vocab_size, 
  29.                                num_layers, 
  30.                                device, 
  31.                                max_len,
  32.                                embed_size, 
  33.                                heads, 
  34.                                forward_expansion, 
  35.                                dropout)
  36.  
  37.         self.src_pad_idx = src_pad_idx
  38.         self.trg_pad_idx = trg_pad_idx
  39.         self.device = device
  40.     
  41.     # 构造mask
  42.     def make_src_mask(self, src):
  43.         # [N,src_len]==>[N,1,1,src_len]
  44.         src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
  45.         return src_mask.to(self.device)
  46.     
  47.     def make_trg_mask(self, trg):
  48.         # 获取目标句子的shape
  49.         N, trg_len = trg.shape
  50.         # 构造mask  [trg_len, trg_len]==>[N, 1, trg_len, trg_len]
  51.         trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand(N, 1, trg_len, trg_len)
  52.         return trg_mask.to(self.device)
  53.     
  54.     # 前向传播
  55.     def forward(self, src, trg):
  56.  
  57.         # 对输入句子构造mask
  58.         src_mask = self.make_src_mask(src)
  59.         # 对目标句子构造mask
  60.         trg_mask = self.make_trg_mask(trg)
  61.  
  62.         # encoder
  63.         enc_src = self.encoder(src, src_mask)
  64.         # decoder
  65.         out = self.decoder(trg, enc_src, src_mask, trg_mask)
  66.         return out

8. 前向传播

下面定义一个输入序列和目标序列，做一次前向传播，验证一下模型是否有误

  1. # -------------------------------------------------- #
  2. #（7）模型测试
  3. # -------------------------------------------------- #
  4. if __name__ == '__main__':
  5.     # 电脑上有GPU就调用它，没有就用CPU
  6.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  7.  
  8.     # 输入
  9.     x = torch.tensor([[1,5,6,4,3,9,5,2,0], [1,8,7,3,4,5,6,7,2]]).to(device)
  10.     # 目标
  11.     trg = torch.tensor([[1,7,4,3,5,9,2,0], [1,5,6,2,4,7,6,2]]).to(device)
  12.     
  13.     src_vocab_size = 10  # 输入句子的长度
  14.     trg_vocab_size = 10  # 目标句子的长度
  15.  
  16.     src_pad_idx = 0  # 对输入句子中的0值做mask
  17.     trg_pad_idx = 0
  18.  
  19.     # 接收模型
  20.     model = Transformer(src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx)
  21.     model = model.to(device)
  22.  
  23.     # 前向传播，参数：输入句子和目标句子
  24.     out = model(x, trg[:,:-1])  # 预测最后一个句子
  25.  
  26.     print(out.shape)  
  27.     # torch.Size([2, 7, 10])
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/345728
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号