bd-Transformer架构详解
Google 2017年论文Attention is all you need提出了Transformer模型,完全基于Attention mechanism,抛弃了传统的CNN和RNN。
1. Transformer架构Transformer解释下这个结构图。首先,Transformer模型也是使用经典的encoder-decoder架构,由encoder和decoder两部分组成。
上图左侧用Nx框出来的,就是我们encoder的一层。encoder一共有6层这样的结构。
上图右侧用Nx框出来的,就是我们decoder的一层。decoder一共有6层这样的结构。
输入序列经过word embedding和positional embedding相加后,输入到encoder中。
输出序列经过word embedding和positional embedding相加后,输入到decoder中。
最后,decoder输出的结果,经过一个线性层,然后计算softmax。
2. Encoderencoder由6层相同的层组成,每一层分别由两部分组成:
第一部分是multi-head self-attention mechanism
第二部分是position-wise feed-forward network,是一个全连接层。
两部分,都有一个残差连接(residual connection),然后接着一个Layer Normalization。
3. Decoder与encoder类似,decoder也是由6个相同层组成,每一个层包括以下3个部分:
第一部分是multi-head self-attention mechanism
第二部分是multi-head context-attention mechanism
第三部分是position-wise feed-forward network
同样,上面三部分中每一部分,都有一个残差连接(residual connection),后接着一个Layer Normalization。
4. Attention机制Attention是指对于某个时刻的输出y,它在输入x上各个部分的注意力。这个注意力可以理解为权重。
attention机制有很多计算方式,下面是一张比较全面的表格:
image.pngseq2seq模型中,使用的是加性注意力(addtion attention)较多。
为什么这种attention叫做addtion attention呢?很简单,对于输入序列隐状态和输出序列的隐状态,它的处理方式很简单,直接合并为
但是transformer模型使用的不是这种attention机制,使用的是另一种,叫做乘性注意力(multiplicative attention)。
那么这种乘性注意力机制是怎么样的呢?从上表中的公式也可以看出来:两个隐状态进行点积!
4.1 Self-attention是什么?上面我们说的attention机制的时候,都会提到两个隐状态,分别是和,前者是输入序列第个位置产生的隐状态,后者是输出序列在第个位置产生的隐状态。
所谓self-attention实际上就是输出序列就是输入序列,因此计算自己的attention得分,就叫做self-attention!
4.2 Context-attention是什么?context-attention是encoder和decoder之间的attention!,所以,也可以成为encoder-decoder attention!
不管是self-attention还是context-attention,它们计算attention分数的时候,可以选择很多方式,比如上面表中提到的:
additive attention
local-base
general
dot-product
scaled dot-product
那么Transformer模型,采用的是哪种呢?答案是:scaled dot-product attention。
4.3 Scaled dot-product attention是什么?论文Attention is all you need里面对于attention机制的描述是这样的:
An attention function can be described as a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility of the query with the corresponding key.
这句话描述得很清楚了。翻译过来就是:通过确定Q和K之间的相似程度来选择V!
用公式来描述更加清晰:
scaled dot-product attention和dot-product attention唯一区别是,scaled dot-product attention有一个缩放因子。
上面公式中表示的是的维度,在论文中,默认是64。
那么为什么需要加上这个缩放因子呢?论文中给出了解释:对于很大时,点积得到的结果维度很大,使得结果处理softmax函数梯度很小的区域。
我们知道,梯度很小时,这对反向传播不利。为了克服这个负面影响,除以一个缩放因子,在一定程度上减缓这种情况。
为什么是呢?论文没有进一步说明。个人觉得你可以使用其他缩放因子,看看模型效果有没有提升。
论文中也提供了一张很清晰的结果图,供大家参考:
image.png首先说明一下我们的是什么:
在encoder的self-attention中,Q、K、V都来自同一个地方(相等),他们是上一层encoder的输出。对于第一层encoder,它们就是word embedding和positional encoding相加得到的输入。
在decoder的self-attention中,Q、K、V都来自同一个地方(相等),他们是上一层decoder的输出。对于第一层decoder,它们就是word embedding和positional encoding相加得到的输入。但是对于decoder,我们不希望它能获得下一个time step,因此我们需要进行sequence masking。
在encoder-decoder attention中,Q来自于decoder的上一层的输出,K和V来自于encoder的输出,K和V是一样的。
三者的维度一样,即。
4.4 Scaled dot-product attention代码实现import numpy as np import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): Scaled dot-product attention mechanism. def __init__(self, attention_dropout=0.0): super(ScaledDotProductAttention, self).__init__() self.dropout = nn.Dropout(attention_dropout) self.softmax = nn.Softmax(dim=2) def forward(self, q, k, v, scale=None, attn_mask=None): 前向传播 args: q: Queries张量,形状[B, L_q, D_q] k: keys张量, 形状[B, L_k, D_k] v: Values张量,形状[B, L_v, D_v] scale: 缩放因子,一个浮点标量 attn_mask: Masking张量,形状[B, L_q, L_k] returns: 上下文张量和attention张量 attention = torch.bmm(q, k.transpose(1, 2)) if scale: attention = attention * scale if attn_mask: # 给需要mask的地方设置一个负无穷 attention = attention.masked_fill_(attn_mask, -np.inf) # 计算softmax attention = self.softmax(attention) # 添加dropout attention = self.dropout(attention) # 和V做点积 context = torch.bmm(attention, v) return context, attention5. Multi-head attention是什么呢?
理解了Scaled dot-product attention,Multi-head attention也很简单了。论文提到,他们发现将Q、K、V通过一个线性映射之后,分成份,对每一份进行scaled dot-product attention效果更好。然后,把各个部分的结果合并起来,再次经过线性映射,得到最终的输出。这就是所谓的multi-head attention。上面的超参数就是heads数量。论文默认是8。
multi-head attention的结构图如下:
image.png
值得注意的是,上面所说的分成份是在维度上面进行切分的。因此,进入到scaled dot-product attention的实际上等于未进入之前的。
Multi-head attention允许模型加入不同位置的表示子空间的信息。
Multi-head attention的公式如下:
其中,
论文中,所以scaled dot-product attention里面的
class MultiHeadAttention(nn.Module): def __init__(self, model_dim=512, num_heads=8, dropout=0.0): super(MultiHeadAttention, self).__init__() self.dim_per_head = model_dim / num_heads self.num_heads = num_heads self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads) self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads) self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads) self.dot_product_attention = ScaledDotProductAttention(dropout) self.linear_final = nn.Linear(model_dim, model_dim) self.dropout = nn.Dropout(dropout) # multi-head attention之后需要做layer norm self.layer_num = nn.LayerNorm(model_dim) def forward(self, query, key, value, attn_mask=None): # 残差连接 residual = query batch_size = key.size(0) # linear projection query = self.linear_q(query) # [B, L, D] key = self.linear_k(key) # [B, L, D] value = self.linear_v(value) # [B, L, D] # split by head query = query.view(batch_size * num_heads, -1, dim_per_head) # [B * 8, , D / 8] key = key.view(batch_size * num_heads, -1, dim_per_head) # value = value.view(batch_size * num_heads, -1, dim_per_head) if attn_mask: attn_mask = attn_mask.repeat(num_heads, 1, 1) # scaled dot product attention scale = (key.size(-1) // num_heads) ** -0.5 context, attention = self.dot_product_attention( query, key, value, scale, attn_mask ) # concat heads context = context.view(batch_size, -1, dim_per_head * num_heads) # final linear projection output = self.linear_final(context) # dropout output = self.dropout(output) # add residual and norm layer output = self.layer_num(residual + output) return output, attention
上面代码中出现了 Residual connection和Layer normalization。下面进行解释:
5.1.1 Residual connection是什么?残差连接其实比较简单!看图就会比较清晰:
image.png假设网络中某个层对输入x作用后的输出为,那么增加residual connection之后,变成:
这个操作被称为shotcut。
残差结构因为增加了一项,该层网络对求偏导时,为常数项1!所以可以在反向传播过程中,梯度连乘,不会造成梯度消失!
5.1.2 Layer normalization是什么?归一化层,主要有这几种方法,BatchNorm(2015年)、LayerNorm(2016年)、InstanceNorm(2016年)、GroupNorm(2018年);
将输入的图像shape记为[N,C,H,W],这几个方法主要区别是:
BatchNorm:batch方向做归一化,计算NHW的均值,对小batchsize效果不好;(BN主要缺点是对batchsize的大小比较敏感,由于每次计算均值和方差是在一个batch上,所以如果batchsize太小,则计算的均值、方差不足以代表整个数据分布)
LayerNorm:channel方向做归一化,计算CHW的均值;(对RNN作用明显)
InstanceNorm:一个batch,一个channel内做归一化。计算HW的均值,用在风格化迁移;(因为在图像风格化中,生成结果主要依赖于某个图像实例,所以对整个batch归一化不适合图像风格化中,因而对HW做归一化。可以加速模型收敛,并且保持每个图像实例之间的独立。)
GroupNorm:将channel方向分group,然后每个group内做归一化,算(C//G)HW的均值;这样与batchsize无关,不受其约束。
Normalization layers6. Mask是什么?mask顾名思义就是掩码,大概意思是对某些值进行掩盖,使其不产生效果.
需要说明的是,Transformer模型中有两种mask。分别是padding mask和sequence mask。其中,padding mask在所有的scaled dot-product attention里都需要用到,而sequence mask只在decoder的self-attention中用到。
所以,我们之前的ScaledDotProductAttention的forward方法里的参数attn_mask在不同的地方有不同的含义。
6.1 Padding mask什么是padding mask呢?回想一下,我们的每个批次输入序列长度是不一样的!也就是说,我们要对输入序列进行对齐!具体来说,就是给较短序列后面填充0。因为这些填充位置,其实没有意义,所以我们的attention机制不应该把注意力放在这些位置上,所以我们需要进行一些处理。
具体做法是:把这些位置的值加上一个非常大的负数(可以是负无穷),这样的话,经过softmax,这些位置的概率就会接近0。
而我们的padding mask实际上是一个张量,每个值都是一个Boolean,值为False的地方就是我们要进行处理的地方。
下面是代码实现:
def padding_mask(seq_q, seq_k): # seq_k和seq_q的形状都是[B,L] len_q = seq_q.size(1) # `PAD` is 0 pad_mask = seq_k.eq(0) pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1) # shape [B,L_q,L_k]
[B,L]- [B,1,L]- [B,L,L]
FFTTFFTTFFTTFFTT6.2 Sequence masksequence mask是为了使得decoder不能看到未来的信息。也就是对于一个序列,在time step为t的时刻,我们的解码输出只能依赖于t时刻之前的输出,而不能依赖t之后的输出。因此我们需要想一个办法,把t之后的信息给隐藏起来。
那具体如何做呢?也很简单:产生一个上三角矩阵,上三角矩阵的值全为1,下三角的值全为0,对角线值也为0。把这个矩阵作用在每一个序列上,就可以达到我们的目的。
具体代码如下:
def sequence_mask(seq): batch_size, seq_len = seq.size() mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8), diagonal=1) mask = mask.unsqueeze(0).expand(batch_size, -1, -1) # [B, L, L] return mask
[B,L,L]
0111001100010000哈佛大学的文章The Annotated Transformer有一张效果图:
值得注意的是,本来mask只需要二维矩阵即可,但是考虑到我们的输入序列都是批量的,所以我们需要把原本二维矩阵扩张成3维张量。上面代码中,已经做了处理。
回到本节开始的问题,attn_mask参数有几种情况?分别是什么意思?
对于decoder的self-attention,里面使用的scaled dot-product attention,同时需要padding mask和sequence mask作为attn_mask,具体实现就是两个mask相加作为attn_mask。
其它情况,attn_mask都等于padding mask。
7. Positional encoding是什么?就目前而言,Transformer架构似乎少了点东西。没错,那就是它对序列的顺序没有约束!我们知道序列的顺序是一个很重要的信息,如果缺失了这个信息,可能我们的结果就是:所有词语都对了,但是无法组成有意义的语句。
为了解决这个问题,论文中提出了positional encoding。一句话概括就是:对序列中的词语出现的位置进行编码!如果对位置进行编码,那么我们的模型就可以捕捉顺序信息。
那么具体怎么做呢?论文的实现是使用正余弦函数。公式如下:
其中,pos是指词语在序列中的位置。可以看出,在偶数位置,使用正弦编码,在奇数位置,使用余弦编码。
上面公式中的是模型的维度,论文默认是512。
这个编码公式的意思就是:给定词语的位置pos,我们可以把它编码成维的向量!也就是说,位置编码的每一个维度对应正弦曲线,波长构成了从到的等比序列。
Postional encoding是对词汇的位置编码。
7.1 Positional encoding代码实现class PositionalEncoding(nn.Module): def __init__(self, d_model, max_seq_len): 初始化 args: d_model: 一个标量。模型的维度,论文默认是512 max_seq_len: 一个标量。文本序列的最大长度 super(PositionalEncoding, self).__init__() # 根据论文给出的公式,构造出PE矩阵 position_encoding = np.array([ [pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)] for pos in range(max_seq_len) ]) # 偶数列使用sin,奇数列使用cos position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2]) position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2]) # 在PE矩阵的一次行,加上一个全是0的向量,代表这`PAD`的positional_encoding # 在word embedding中也会经常加上`UNK`,代表位置单词的word embedding,两者十分类似 # 那么为什么需要这个额外的PAD的编码呢?很简单,因为文本序列的长度不易,我们需要对齐, # 短的序列我们使用0在结尾不全,我们也需要这些补全位置的编码,也就是`PAD`对应的位置编码 pad_row = torch.zeros([1, d_model]) position_encoding = torch.cat((pad_row, position_encoding)) # 嵌入操作,+1是因为增加了`PAD`这个补全位置的编码 # word embedding中如果词典增加`UNK`,我们也需要+1。 self.position_encoding = nn.Embedding(max_seq_len+1, d_model) self.position_encoding.weight = nn.Parameter(position_encoding, requires_grad=False) def forward(self, input_len): 神经网络前向传播 args: input_len: 一个张量,形状为[BATCH_SIZE, 1]。每一个张量的值代表这一批文本序列中对应的长度。 returns: 返回这一批序列的位置编码,进行了对齐。 # 找出这一批序列的最大长度 max_len = torch.max(input_len) # 对每一个序列的位置进行对齐,在原序列位置的后面补上0 # 这里range从1开始也是因为要避开PAD(0)的位置 input_pos = torch.LongTensor( [list(range(1, len+1)) + [0] * (max_len-len) for len in input_len] ) return self.position_encoding(input_pos)8. Word embedding是什么?
Word embedding是对序列中的词汇的编码,把每一个词汇编码成维的向量!它实际上就是一个二维浮点矩阵,里面的权重是可训练参数,我们只需要把这个矩阵构建出来就完成了word embedding的工作。
embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0)
上面vocab_size是词典大小,embedding_size是词嵌入的维度大小,论文里面就是等于。所以word embedding矩阵就是一个vocab_size*embedding_size的二维张量。
9. Position-wise Feed-Forward netword是什么?这是一个全连接网络,包含连个线性变换和一个非线性函数(ReLU)。公式如下:
这个线性变换在不同的位置都是一样的,并且在不同的层之间使用不同的参数。
论文提到,这个公式还可以用两个核大小为1的一维卷积来解释,卷积的输入输出都是,中间层维度是。
代码如下:
class PositionalWiseFeedForward(nn.Module): def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0): super(PositionalWiseFeedForward, self).__init__() self.w1 = nn.Conv1d(model_dim, ffn_dim, 1) self.w2 = nn.Conv2d(model_dim, ffn_dim, 1) self.dropout = nn.Dropout(dropout) self.layer_norm = nn.LayerNorm(model_dim) def forward(self, x): output = x.transpose(1, 2) output = self.w2(F.relu(self.w1(output))) output = self.dropout(output.transpose(1, 2)) # add residual and norm layer output = self.layer_norm(x + output) return output10. 完整代码
至此,所有的细节都解释完了。
import numpy as np import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): Scaled dot-product attention mechanism. def __init__(self, attention_dropout=0.0): super(ScaledDotProductAttention, self).__init__() self.dropout = nn.Dropout(attention_dropout) self.softmax = nn.Softmax(dim=2) def forward(self, q, k, v, scale=None, attn_mask=None): 前向传播 args: q: Queries张量,形状[B, L_q, D_q] k: keys张量, 形状[B, L_k, D_k] v: Values张量,形状[B, L_v, D_v] scale: 缩放因子,一个浮点标量 attn_mask: Masking张量,形状[B, L_q, L_k] returns: 上下文张量和attention张量 attention = torch.bmm(q, k.transpose(1, 2)) if scale: attention = attention * scale if attn_mask: # 给需要mask的地方设置一个负无穷 attention = attention.masked_fill_(attn_mask, -np.inf) # 计算softmax attention = self.softmax(attention) # 添加dropout attention = self.dropout(attention) # 和V做点积 context = torch.bmm(attention, v) return context, attention class MultiHeadAttention(nn.Module): def __init__(self, model_dim=512, num_heads=8, dropout=0.0): super(MultiHeadAttention, self).__init__() self.dim_per_head = model_dim / num_heads self.num_heads = num_heads self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads) self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads) self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads) self.dot_product_attention = ScaledDotProductAttention(dropout) self.linear_final = nn.Linear(model_dim, model_dim) self.dropout = nn.Dropout(dropout) # multi-head attention之后需要做layer norm self.layer_num = nn.LayerNorm(model_dim) def forward(self, query, key, value, attn_mask=None): # 残差连接 residual = query batch_size = key.size(0) # linear projection query = self.linear_q(query) # [B, L, D] key = self.linear_k(key) # [B, L, D] value = self.linear_v(value) # [B, L, D] # split by head query = query.view(batch_size * num_heads, -1, dim_per_head) # [B * 8, , D / 8] key = key.view(batch_size * num_heads, -1, dim_per_head) # value = value.view(batch_size * num_heads, -1, dim_per_head) if attn_mask: attn_mask = attn_mask.repeat(num_heads, 1, 1) # scaled dot product attention scale = (key.size(-1) // num_heads) ** -0.5 context, attention = self.dot_product_attention( query, key, value, scale, attn_mask ) # concat heads context = context.view(batch_size, -1, dim_per_head * num_heads) # final linear projection output = self.linear_final(context) # dropout output = self.dropout(output) # add residual and norm layer output = self.layer_num(residual + output) return output, attention def padding_mask(seq_q, seq_k): # seq_k和seq_q的形状都是[B,L] len_q = seq_q.size(1) # `PAD` is 0 pad_mask = seq_k.eq(0) pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1) # shape [B,L_q,L_k] def sequence_mask(seq): batch_size, seq_len = seq.size() mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8), diagonal=1) mask = mask.unsqueeze(0).expand(batch_size, -1, -1) # [B, L, L] return mask class PositionalEncoding(nn.Module): def __init__(self, d_model, max_seq_len): 初始化 args: d_model: 一个标量。模型的维度,论文默认是512 max_seq_len: 一个标量。文本序列的最大长度 super(PositionalEncoding, self).__init__() # 根据论文给出的公式,构造出PE矩阵 position_encoding = np.array([ [pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)] for pos in range(max_seq_len) ]) # 偶数列使用sin,奇数列使用cos position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2]) position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2]) # 在PE矩阵的一次行,加上一个全是0的向量,代表这`PAD`的positional_encoding # 在word embedding中也会经常加上`UNK`,代表位置单词的word embedding,两者十分类似 # 那么为什么需要这个额外的PAD的编码呢?很简单,因为文本序列的长度不易,我们需要对齐, # 短的序列我们使用0在结尾不全,我们也需要这些补全位置的编码,也就是`PAD`对应的位置编码 pad_row = torch.zeros([1, d_model]) position_encoding = torch.cat((pad_row, position_encoding)) # 嵌入操作,+1是因为增加了`PAD`这个补全位置的编码 # word embedding中如果词典增加`UNK`,我们也需要+1。 self.position_encoding = nn.Embedding(max_seq_len+1, d_model) self.position_encoding.weight = nn.Parameter(position_encoding, requires_grad=False) def forward(self, input_len): 神经网络前向传播 args: input_len: 一个张量,形状为[BATCH_SIZE, 1]。每一个张量的值代表这一批文本序列中对应的长度。 returns: 返回这一批序列的位置编码,进行了对齐。 # 找出这一批序列的最大长度 max_len = torch.max(input_len) # 对每一个序列的位置进行对齐,在原序列位置的后面补上0 # 这里range从1开始也是因为要避开PAD(0)的位置 input_pos = torch.LongTensor( [list(range(1, len+1)) + [0] * (max_len-len) for len in input_len] ) return self.position_encoding(input_pos) # embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0) # 获得输入的词嵌入编码 # seq_embedding = seq_embedding(inputs) * np.sqrt(d_model) class PositionalWiseFeedForward(nn.Module): def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0): super(PositionalWiseFeedForward, self).__init__() self.w1 = nn.Conv1d(model_dim, ffn_dim, 1) self.w2 = nn.Conv2d(model_dim, ffn_dim, 1) self.dropout = nn.Dropout(dropout) self.layer_norm = nn.LayerNorm(model_dim) def forward(self, x): output = x.transpose(1, 2) output = self.w2(F.relu(self.w1(output))) output = self.dropout(output.transpose(1, 2)) # add residual and norm layer output = self.layer_norm(x + output) return output class EncoderLayer(nn.Module): Encoder的一层。 def __init__(self, model_dim=512, num_heads=8, ffn_dim=2048, dropout=0.0): super(EncoderLayer, self).__init__() self.attention = MultiHeadAttention(model_dim, num_heads, dropout) self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout) def forward(self, inputs, attn_mask=None): # self attention context, attention = self.attention(inputs, inputs, inputs, attn_mask) # feed forward network output = self.feed_forward(context) return output, attention
self.encoder_layers = nn.ModuleList( [EncoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in range(num_layers)] ) self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0) self.pos_embedding = PositionalEncoding(model_dim, max_seq_len) def forward(self, inputs, inputs_len): output = self.seq_embedding(inputs) output += self.pos_embedding(inputs_len) self_attention_mask = padding_mask(inputs, inputs) attentions = [] for encoder in self.encoder_layers: output, attention = encoder(output, self_attention_mask) attentions.append(attention) return output, attentions class DecoderLayer(nn.Module): def __init__(self, model_dim, num_heads=8, ffn_dim=2048, dropout=0.0): super(DecoderLayer, self).__init__() self.attention = MultiHeadAttention(model_dim, num_heads, dropout) self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout) def forward(self, dec_inputs, enc_outputs, self_attn_mask=None, context_attn_mask=None): # self attention, all inputs are decoder inputs dec_output, self_attention = self.attention(dec_inputs, dec_inputs, dec_inputs, self_attn_mask) # context attention # query is decoder s outputs, key and value are encoder s inputs dec_output, context_attention = self.attention(dec_output, enc_outputs, enc_outputs, context_attn_mask) # decoder s output, or context dec_output = self.feed_forward(dec_output) return dec_output, self_attention, context_attention class Decoder(nn.Module): def __init__(self, vocab_size, max_seq_len, num_layers=6, model_dim=512, num_heads=8, ffn_dim=2048, dropout=0.0): super(Decoder).__init__() self.num_layers = num_layers self.decoder_layers = nn.ModuleList( [DecoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in range(num_layers)] ) self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0) self.pos_embedding = PositionalEncoding(model_dim, max_seq_len) def forward(self, inputs, inputs_len, enc_output, context_attn_mask=None): output = self.seq_embedding(inputs) output += self.pos_embedding(inputs_len) self_attention_padding_mask = padding_mask(inputs, inputs) seq_mask = sequence_mask(inputs) self_attn_mask = torch.gt((self_attention_padding_mask + seq_mask), 0) self_attentions = [] context_attentions = [] for decoder in self.decoder_layers: output, self_attn, context_attn = decoder( output, enc_output, self_attn_mask, context_attn_mask) self_attentions.append(self_attn) context_attentions.append(context_attn) return output, self_attentions, context_attentions class Transformer(nn.Module): def __init__(self, src_vocab_size, src_max_len, tgt_vocab_size, tgt_max_len, num_layers=6, model_dim=512, num_heads=8, ffn_dim=2048, dropout=0.0): super(Transformer).__init__() self.encoder = Encoder(src_vocab_size, src_max_len, num_layers, model_dim, num_heads, ffn_dim, dropout) self.decoder = Decoder(tgt_vocab_size, tgt_max_len, num_layers, model_dim, num_heads, ffn_dim, dropout) self.linear = nn.Linear(model_dim, tgt_vocab_size, bias=False) self.softmax = nn.Softmax() def forward(self, src_seq, src_len, tgt_seq, tgt_len): context_attn_mask = padding_mask(tgt_seq, src_seq) output, enc_self_attn = self.encoder(src_seq, src_len) output, dec_self_attn, ctx_attn = self.decoder(tgt_seq, tgt_len, output, context_attn_mask) output = self.linear(output) output = self.softmax(output) return output, enc_self_attn, dec_self_attn, ctx_attn
本站是提供个人知识管理的网络存储空间,所有内容均由用户发布,不代表本站观点。请注意甄别内容中的联系方式、诱导购买等信息,谨防诈骗。如发现有害或侵权内容,请点击一键举报。
0条评论