• 150632

    文章

  • 1244

    评论

  • 13

    友链

  • 最近新加了换肤功能,大家多来逛逛吧~~~~
  • 喜欢这个网站的朋友可以加一下QQ群,我们一起交流技术。

熬了一晚上,我从零实现了Transformer模型,把代码讲给你听


 

自从彻底搞懂 Self_Attention 机制之后,笔者对 Transformer 模型的理解直接从地下一层上升到大气层,任督二脉呼之欲出。夜夜入睡之前,那句柔情百转的"Attention is all you need"时常在耳畔环绕,情到深处不禁拍床叫好。于是在肾上腺素的驱使下,笔者熬了一个晚上,终于实现了 Transformer 模型。

1. 模型总览

代码讲解之前,首先放出这张经典的模型架构图。下面的内容中,我会将每个模块的实现思路以及笔者在Coding过程中的感悟知无不答。没有代码基础的读者不要慌张,笔者也是最近才入门的,所写Pytorch代码没有花里胡哨,所用变量名词尽量保持与论文一致,对新手十分友好。

我们观察模型的结构图,Transformer模型包含哪些模块?笔者将其分为以下几个部分:

接下来我们首先逐个讲解,最后将其拼接完成模型的复现。

2. config

下面是这个Demo所用的库文件以及一些超参的信息。 单独实现一个Config类保存的原因是,方便日后复用。直接将模型部分复制,所用超参保存在新项目的Config类中即可 。这里不过多赘述。

  
  1. import torch 
  2. import torch.nn as nn 
  3. import numpy as np 
  4. import math 
  5.  
  6.  
  7. class Config(object): 
  8.     def __init__(self): 
  9.         self.vocab_size = 6 
  10.  
  11.         self.d_model = 20 
  12.         self.n_heads = 2 
  13.  
  14.         assert self.d_model % self.n_heads == 0 
  15.         dim_k  = d_model % n_heads 
  16.         dim_v = d_model % n_heads 
  17.  
  18.         self.padding_size = 30 
  19.         self.UNK = 5 
  20.         self.PAD = 4 
  21.  
  22.         self.N = 6 
  23.         self.p = 0.1 
  24.  
  25. config = Config() 

3. Embedding

Embedding部分接受原始的文本输入(batch_size*seq_len,例:[[1,3,10,5],[3,4,5],[5,3,1,1]]),叠加一个普通的Embedding层以及一个Positional Embedding层,输出最后结果。

在这一层中,输入的是一个list: [batch_size * seq_len],输出的是一个tensor:[batch_size * seq_len * d_model]

普通的 Embedding 层想说两点:

  •  
  •    
    1. torch.nn.Embedding 
    2. torch.nn.Embedding 
    3. padding_idx 
  • 在padding过程中,短补长截

  
  1. class Embedding(nn.Module): 
  2.     def __init__(self,vocab_size): 
  3.         super(Embedding, self).__init__() 
  4.         # 一个普通的 embedding层,我们可以通过设置padding_idx=config.PAD 来实现论文中的 padding_mask 
  5.         self.embedding = nn.Embedding(vocab_size,config.d_model,padding_idx=config.PAD) 
  6.  
  7.  
  8.     def forward(self,x): 
  9.         # 根据每个句子的长度,进行padding,短补长截 
  10.         for i in range(len(x)): 
  11.             if len(x[i]) < config.padding_size: 
  12.                 x[i].extend([config.UNK] * (config.padding_size - len(x[i]))) # 注意 UNK是你词表中用来表示oov的token索引,这里进行了简化,直接假设为6 
  13.             else
  14.                 x[i] = x[i][:config.padding_size] 
  15.         x = self.embedding(torch.tensor(x)) # batch_size * seq_len * d_model 
  16.         return x 

关于Positional Embedding,我们需要参考论文给出的公式。说一句题外话,在作者的实验中对比了Positional Embedding与单独采用一个Embedding训练模型对位置的感知两种方式,模型效果相差无几。

  
  1. class Positional_Encoding(nn.Module): 
  2.  
  3.     def __init__(self,d_model): 
  4.         super(Positional_Encoding,self).__init__() 
  5.         self.d_model = d_model 
  6.  
  7.  
  8.     def forward(self,seq_len,embedding_dim): 
  9.         positional_encoding = np.zeros((seq_len,embedding_dim)) 
  10.         for pos in range(positional_encoding.shape[0]): 
  11.             for i in range(positional_encoding.shape[1]): 
  12.                 positional_encoding[pos][i] = math.sin(pos/(10000**(2*i/self.d_model))) if i % 2 == 0 else math.cos(pos/(10000**(2*i/self.d_model))) 
  13.         return torch.from_numpy(positional_encoding) 

4. Encoder

Muti_head_Attention

这一部分是模型的核心内容,理论部分就不过多讲解了,读者可以参考文章开头的第一个传送门,文中有基础的代码实现。

Encoder 中的 Muti_head_Attention 不需要Mask,因此与我们上一篇文章中的实现方式相同。

为了避免模型信息泄露的问题,Decoder 中的 Muti_head_Attention 需要Mask。这一节中我们重点讲解Muti_head_Attention中Mask机制的实现。

如果读者阅读了我们的上一篇文章,可以发现下面的代码有一点小小的不同,主要体现在 forward 函数的参数。

  •  
  • requires_mask:是否采用Mask机制,在Decoder中设置为True

  
  1. class Mutihead_Attention(nn.Module): 
  2.     def __init__(self,d_model,dim_k,dim_v,n_heads): 
  3.         super(Mutihead_Attention, self).__init__() 
  4.         self.dim_v = dim_v 
  5.         self.dim_k = dim_k 
  6.         self.n_heads = n_heads 
  7.  
  8.         self.q = nn.Linear(d_model,dim_k) 
  9.         self.k = nn.Linear(d_model,dim_k) 
  10.         self.v = nn.Linear(d_model,dim_v) 
  11.  
  12.         self.o = nn.Linear(dim_v,d_model) 
  13.         self.norm_fact = 1 / math.sqrt(d_model) 
  14.  
  15.     def generate_mask(self,dim): 
  16.         # 此处是 sequence mask ,防止 decoder窥视后面时间步的信息。 
  17.         # padding mask 在数据输入模型之前完成。 
  18.         matirx = np.ones((dim,dim)) 
  19.         mask = torch.Tensor(np.tril(matirx)) 
  20.  
  21.         return mask==1 
  22.  
  23.     def forward(self,x,y,requires_mask=False): 
  24.         assert self.dim_k % self.n_heads == 0 and self.dim_v % self.n_heads == 0 
  25.         # size of x : [batch_size * seq_len * batch_size] 
  26.         # 对 x 进行自注意力 
  27.         Q = self.q(x).reshape(-1,x.shape[0],x.shape[1],self.dim_k // self.n_heads) # n_heads * batch_size * seq_len * dim_k 
  28.         K = self.k(x).reshape(-1,x.shape[0],x.shape[1],self.dim_k // self.n_heads) # n_heads * batch_size * seq_len * dim_k 
  29.         V = self.v(y).reshape(-1,y.shape[0],y.shape[1],self.dim_v // self.n_heads) # n_heads * batch_size * seq_len * dim_v 
  30.         # print("Attention V shape : {}".format(V.shape)) 
  31.         attention_score = torch.matmul(Q,K.permute(0,1,3,2)) * self.norm_fact 
  32.         if requires_mask: 
  33.             mask = self.generate_mask(x.shape[1]) 
  34.             attention_score.masked_fill(mask,value=float("-inf")) # 注意这里的小Trick,不需要将Q,K,V 分别MASK,只MASKSoftmax之前的结果就好了 
  35.         output = torch.matmul(attention_score,V).reshape(y.shape[0],y.shape[1],-1
  36.         # print("Attention output shape : {}".format(output.shape)) 
  37.  
  38.         output = self.o(output) 
  39.         return output 

Feed Forward

这一部分实现很简单,两个Linear中连接Relu即可,目的是为模型增添非线性信息,提高模型的拟合能力。

  
  1. class Feed_Forward(nn.Module): 
  2.     def __init__(self,input_dim,hidden_dim=2048): 
  3.         super(Feed_Forward, self).__init__() 
  4.         self.L1 = nn.Linear(input_dim,hidden_dim) 
  5.         self.L2 = nn.Linear(hidden_dim,input_dim) 
  6.  
  7.     def forward(self,x): 
  8.         output = nn.ReLU()(self.L1(x)) 
  9.         output = self.L2(output) 
  10.         return output 

Add & LayerNorm

这一节我们实现论文中提出的残差连接以及LayerNorm。

论文中关于这部分给出公式:

代码中的dropout,在论文中也有所解释,对输入layer_norm的tensor进行dropout,对模型的性能影响还是蛮大的。

代码中的参数 sub_layer ,可以是Feed Forward,也可以是Muti_head_Attention。

  
  1. class Add_Norm(nn.Module): 
  2.     def __init__(self): 
  3.         self.dropout = nn.Dropout(config.p) 
  4.         super(Add_Norm, self).__init__() 
  5.  
  6.     def forward(self,x,sub_layer,**kwargs): 
  7.         sub_output = sub_layer(x,**kwargs) 
  8.         # print("{} output : {}".format(sub_layer,sub_output.size())) 
  9.         x = self.dropout(x + sub_output) 
  10.  
  11.         layer_norm = nn.LayerNorm(x.size()[1:]) 
  12.         out = layer_norm(x) 
  13.         return out 

OK,Encoder中所有模块我们已经讲解完毕,接下来我们将其拼接作为Encoder

  
  1. class Encoder(nn.Module): 
  2.     def __init__(self): 
  3.         super(Encoder, self).__init__() 
  4.         self.positional_encoding = Positional_Encoding(config.d_model) 
  5.         self.muti_atten = Mutihead_Attention(config.d_model,config.dim_k,config.dim_v,config.n_heads) 
  6.         self.feed_forward = Feed_Forward(config.d_model) 
  7.  
  8.         self.add_norm = Add_Norm() 
  9.  
  10.  
  11.     def forward(self,x): # batch_size * seq_len 并且 x 的类型不是tensor,是普通list 
  12.  
  13.         x += self.positional_encoding(x.shape[1],config.d_model) 
  14.         # print("After positional_encoding: {}".format(x.size())) 
  15.         output = self.add_norm(x,self.muti_atten,y=x) 
  16.         output = self.add_norm(output,self.feed_forward) 
  17.  
  18.         return output 

5.Decoder

在 Encoder 部分的讲解中,我们已经实现了大部分Decoder的模块。Decoder的Muti_head_Attention引入了Mask机制,Decoder与Encoder 中模块的拼接方式不同。以上两点读者在Coding的时候需要注意。

  
  1. class Decoder(nn.Module): 
  2.     def __init__(self): 
  3.         super(Decoder, self).__init__() 
  4.         self.positional_encoding = Positional_Encoding(config.d_model) 
  5.         self.muti_atten = Mutihead_Attention(config.d_model,config.dim_k,config.dim_v,config.n_heads) 
  6.         self.feed_forward = Feed_Forward(config.d_model) 
  7.         self.add_norm = Add_Norm() 
  8.  
  9.     def forward(self,x,encoder_output): # batch_size * seq_len 并且 x 的类型不是tensor,是普通list 
  10.         # print(x.size()) 
  11.         x += self.positional_encoding(x.shape[1],config.d_model) 
  12.         # print(x.size()) 
  13.         # 第一个 sub_layer 
  14.         output = self.add_norm(x,self.muti_atten,y=x,requires_mask=True) 
  15.         # 第二个 sub_layer 
  16.         output = self.add_norm(output,self.muti_atten,y=encoder_output,requires_mask=True) 
  17.         # 第三个 sub_layer 
  18.         output = self.add_norm(output,self.feed_forward) 
  19.  
  20.  
  21.         return output 

6.Transformer

至此,所有内容已经铺垫完毕,我们开始组装Transformer模型。论文中提到,Transformer中堆叠了6个我们上文中实现的Encoder 和 Decoder。这里笔者采用 nn.Sequential 实现了堆叠操作。

Output模块的 Linear 和 Softmax 的实现也包含在下面的代码中

  
  1. class Transformer_layer(nn.Module): 
  2.     def __init__(self): 
  3.         super(Transformer_layer, self).__init__() 
  4.         self.encoder = Encoder() 
  5.         self.decoder = Decoder() 
  6.  
  7.     def forward(self,x): 
  8.         x_input,x_output = x 
  9.         encoder_output = self.encoder(x_input) 
  10.         decoder_output = self.decoder(x_output,encoder_output) 
  11.         return (encoder_output,decoder_output) 
  12.  
  13.  
  14. class Transformer(nn.Module): 
  15.     def __init__(self,N,vocab_size,output_dim): 
  16.         super(Transformer, self).__init__() 
  17.         self.embedding_input = Embedding(vocab_size=vocab_size) 
  18.         self.embedding_output = Embedding(vocab_size=vocab_size) 
  19.  
  20.         self.output_dim = output_dim 
  21.         self.linear = nn.Linear(config.d_model,output_dim) 
  22.         self.softmax = nn.Softmax(dim=-1
  23.         self.model = nn.Sequential(*[Transformer_layer() for _ in range(N)]) 
  24.  
  25.  
  26.     def forward(self,x): 
  27.         x_input , x_output = x 
  28.         x_input = self.embedding_input(x_input) 
  29.         x_output = self.embedding_output(x_output) 
  30.  
  31.         _ , output = self.model((x_input,x_output)) 
  32.  
  33.         output = self.linear(output) 
  34.         output = self.softmax(output) 
  35.  
  36.         return output 

完整代码

  
  1. @Author:Yifx 
  2. @Contact: Xxuyifan1999@163.com 
  3. @Time:2021/9/16 20:02 
  4. @Software: PyCharm 
  5.  
  6. ""
  7. 文件说明: 
  8. ""
  9.  
  10. import torch 
  11. import torch.nn as nn 
  12. import numpy as np 
  13. import math 
  14.  
  15.  
  16. class Config(object): 
  17.     def __init__(self): 
  18.         self.vocab_size = 6 
  19.  
  20.         self.d_model = 20 
  21.         self.n_heads = 2 
  22.  
  23.         assert self.d_model % self.n_heads == 0 
  24.         dim_k  = self.d_model // self.n_heads 
  25.         dim_v = self.d_model // self.n_heads 
  26.  
  27.  
  28.  
  29.         self.padding_size = 30 
  30.         self.UNK = 5 
  31.         self.PAD = 4 
  32.  
  33.         self.N = 6 
  34.         self.p = 0.1 
  35.  
  36. config = Config() 
  37.  
  38.  
  39. class Embedding(nn.Module): 
  40.     def __init__(self,vocab_size): 
  41.         super(Embedding, self).__init__() 
  42.         # 一个普通的 embedding层,我们可以通过设置padding_idx=config.PAD 来实现论文中的 padding_mask 
  43.         self.embedding = nn.Embedding(vocab_size,config.d_model,padding_idx=config.PAD) 
  44.  
  45.  
  46.     def forward(self,x): 
  47.         # 根据每个句子的长度,进行padding,短补长截 
  48.         for i in range(len(x)): 
  49.             if len(x[i]) < config.padding_size: 
  50.                 x[i].extend([config.UNK] * (config.padding_size - len(x[i]))) # 注意 UNK是你词表中用来表示oov的token索引,这里进行了简化,直接假设为6 
  51.             else
  52.                 x[i] = x[i][:config.padding_size] 
  53.         x = self.embedding(torch.tensor(x)) # batch_size * seq_len * d_model 
  54.         return x 
  55.  
  56.  
  57.  
  58. class Positional_Encoding(nn.Module): 
  59.  
  60.     def __init__(self,d_model): 
  61.         super(Positional_Encoding,self).__init__() 
  62.         self.d_model = d_model 
  63.  
  64.  
  65.     def forward(self,seq_len,embedding_dim): 
  66.         positional_encoding = np.zeros((seq_len,embedding_dim)) 
  67.         for pos in range(positional_encoding.shape[0]): 
  68.             for i in range(positional_encoding.shape[1]): 
  69.                 positional_encoding[pos][i] = math.sin(pos/(10000**(2*i/self.d_model))) if i % 2 == 0 else math.cos(pos/(10000**(2*i/self.d_model))) 
  70.         return torch.from_numpy(positional_encoding) 
  71.  
  72.  
  73. class Mutihead_Attention(nn.Module): 
  74.     def __init__(self,d_model,dim_k,dim_v,n_heads): 
  75.         super(Mutihead_Attention, self).__init__() 
  76.         self.dim_v = dim_v 
  77.         self.dim_k = dim_k 
  78.         self.n_heads = n_heads 
  79.  
  80.         self.q = nn.Linear(d_model,dim_k) 
  81.         self.k = nn.Linear(d_model,dim_k) 
  82.         self.v = nn.Linear(d_model,dim_v) 
  83.  
  84.         self.o = nn.Linear(dim_v,d_model) 
  85.         self.norm_fact = 1 / math.sqrt(d_model) 
  86.  
  87.     def generate_mask(self,dim): 
  88.         # 此处是 sequence mask ,防止 decoder窥视后面时间步的信息。 
  89.         # padding mask 在数据输入模型之前完成。 
  90.         matirx = np.ones((dim,dim)) 
  91.         mask = torch.Tensor(np.tril(matirx)) 
  92.  
  93.         return mask==1 
  94.  
  95.     def forward(self,x,y,requires_mask=False): 
  96.         assert self.dim_k % self.n_heads == 0 and self.dim_v % self.n_heads == 0 
  97.         # size of x : [batch_size * seq_len * batch_size] 
  98.         # 对 x 进行自注意力 
  99.         Q = self.q(x).reshape(-1,x.shape[0],x.shape[1],self.dim_k // self.n_heads) # n_heads * batch_size * seq_len * dim_k 
  100.         K = self.k(x).reshape(-1,x.shape[0],x.shape[1],self.dim_k // self.n_heads) # n_heads * batch_size * seq_len * dim_k 
  101.         V = self.v(y).reshape(-1,y.shape[0],y.shape[1],self.dim_v // self.n_heads) # n_heads * batch_size * seq_len * dim_v 
  102.         # print("Attention V shape : {}".format(V.shape)) 
  103.         attention_score = torch.matmul(Q,K.permute(0,1,3,2)) * self.norm_fact 
  104.         if requires_mask: 
  105.             mask = self.generate_mask(x.shape[1]) 
  106.             # masked_fill 函数中,对Mask位置为True的部分进行Mask 
  107.             attention_score.masked_fill(mask,value=float("-inf")) # 注意这里的小Trick,不需要将Q,K,V 分别MASK,只MASKSoftmax之前的结果就好了 
  108.         output = torch.matmul(attention_score,V).reshape(y.shape[0],y.shape[1],-1
  109.         # print("Attention output shape : {}".format(output.shape)) 
  110.  
  111.         output = self.o(output) 
  112.         return output 
  113.  
  114.  
  115. class Feed_Forward(nn.Module): 
  116.     def __init__(self,input_dim,hidden_dim=2048): 
  117.         super(Feed_Forward, self).__init__() 
  118.         self.L1 = nn.Linear(input_dim,hidden_dim) 
  119.         self.L2 = nn.Linear(hidden_dim,input_dim) 
  120.  
  121.     def forward(self,x): 
  122.         output = nn.ReLU()(self.L1(x)) 
  123.         output = self.L2(output) 
  124.         return output 
  125.  
  126. class Add_Norm(nn.Module): 
  127.     def __init__(self): 
  128.         self.dropout = nn.Dropout(config.p) 
  129.         super(Add_Norm, self).__init__() 
  130.  
  131.     def forward(self,x,sub_layer,**kwargs): 
  132.         sub_output = sub_layer(x,**kwargs) 
  133.         # print("{} output : {}".format(sub_layer,sub_output.size())) 
  134.         x = self.dropout(x + sub_output) 
  135.  
  136.         layer_norm = nn.LayerNorm(x.size()[1:]) 
  137.         out = layer_norm(x) 
  138.         return out 
  139.  
  140.  
  141. class Encoder(nn.Module): 
  142.     def __init__(self): 
  143.         super(Encoder, self).__init__() 
  144.         self.positional_encoding = Positional_Encoding(config.d_model) 
  145.         self.muti_atten = Mutihead_Attention(config.d_model,config.dim_k,config.dim_v,config.n_heads) 
  146.         self.feed_forward = Feed_Forward(config.d_model) 
  147.  
  148.         self.add_norm = Add_Norm() 
  149.  
  150.  
  151.     def forward(self,x): # batch_size * seq_len 并且 x 的类型不是tensor,是普通list 
  152.  
  153.         x += self.positional_encoding(x.shape[1],config.d_model) 
  154.         # print("After positional_encoding: {}".format(x.size())) 
  155.         output = self.add_norm(x,self.muti_atten,y=x) 
  156.         output = self.add_norm(output,self.feed_forward) 
  157.  
  158.         return output 
  159.  
  160. # 在 Decoder 中,Encoder的输出作为Query和KEy输出的那个东西。即 Decoder的Input作为V。此时是可行的 
  161. # 因为在输入过程中,我们有一个padding操作,将Inputs和Outputs的seq_len这个维度都拉成一样的了 
  162. # 我们知道,QK那个过程得到的结果是 batch_size * seq_len * seq_len .既然 seq_len 一样,那么我们可以这样操作 
  163. # 这样操作的意义是,Outputs 中的 token 分别对于 Inputs 中的每个token作注意力 
  164.  
  165. class Decoder(nn.Module): 
  166.     def __init__(self): 
  167.         super(Decoder, self).__init__() 
  168.         self.positional_encoding = Positional_Encoding(config.d_model) 
  169.         self.muti_atten = Mutihead_Attention(config.d_model,config.dim_k,config.dim_v,config.n_heads) 
  170.         self.feed_forward = Feed_Forward(config.d_model) 
  171.         self.add_norm = Add_Norm() 
  172.  
  173.     def forward(self,x,encoder_output): # batch_size * seq_len 并且 x 的类型不是tensor,是普通list 
  174.         # print(x.size()) 
  175.         x += self.positional_encoding(x.shape[1],config.d_model) 
  176.         # print(x.size()) 
  177.         # 第一个 sub_layer 
  178.         output = self.add_norm(x,self.muti_atten,y=x,requires_mask=True) 
  179.         # 第二个 sub_layer 
  180.         output = self.add_norm(x,self.muti_atten,y=encoder_output,requires_mask=True) 
  181.         # 第三个 sub_layer 
  182.         output = self.add_norm(output,self.feed_forward) 
  183.         return output 
  184.  
  185. class Transformer_layer(nn.Module): 
  186.     def __init__(self): 
  187.         super(Transformer_layer, self).__init__() 
  188.         self.encoder = Encoder() 
  189.         self.decoder = Decoder() 
  190.  
  191.     def forward(self,x): 
  192.         x_input,x_output = x 
  193.         encoder_output = self.encoder(x_input) 
  194.         decoder_output = self.decoder(x_output,encoder_output) 
  195.         return (encoder_output,decoder_output) 
  196.  
  197. class Transformer(nn.Module): 
  198.     def __init__(self,N,vocab_size,output_dim): 
  199.         super(Transformer, self).__init__() 
  200.         self.embedding_input = Embedding(vocab_size=vocab_size) 
  201.         self.embedding_output = Embedding(vocab_size=vocab_size) 
  202.  
  203.         self.output_dim = output_dim 
  204.         self.linear = nn.Linear(config.d_model,output_dim) 
  205.         self.softmax = nn.Softmax(dim=-1
  206.         self.model = nn.Sequential(*[Transformer_layer() for _ in range(N)]) 
  207.  
  208.  
  209.     def forward(self,x): 
  210.         x_input , x_output = x 
  211.         x_input = self.embedding_input(x_input) 
  212.         x_output = self.embedding_output(x_output) 
  213.  
  214.         _ , output = self.model((x_input,x_output)) 
  215.  
  216.         output = self.linear(output) 
  217.         output = self.softmax(output) 
  218.  
  219.         return output 

【编辑推荐】



695856371Web网页设计师②群 | 喜欢本站的朋友可以收藏本站,或者加入我们大家一起来交流技术!

0条评论

Loading...


发表评论

电子邮件地址不会被公开。 必填项已用*标注

自定义皮肤 主体内容背景
打开支付宝扫码付款购买视频教程
遇到问题联系客服QQ:419400980
注册梁钟霖个人博客