10.4 Transformer实现#

在前面两节内容中我们分别详细介绍了Transformer模型的原理与多头注意力机制的实现过程,接下来,我们将会一步一步地来详细介绍如何通过 PyTorch框架实现Transformer的整体网络结构, 包括嵌入层、编码器和解码器等等。下面,首先介绍嵌入层的实现过程。

10.4.1 嵌入层实现#

1. Token Embedding 实现

这里首先要实现的是最基础的字符嵌入层,通过PyTorch框架只需要一行代码便可以得到一个字符嵌入层,如下所示:

1 class TokenEmbedding(nn.Module):
2     def __init__(self, vocab_size, emb_size):
3         super(TokenEmbedding, self).__init__()
4         self.embedding = nn.Embedding(vocab_size, emb_size)
5         self.emb_size = emb_size
6     def forward(self, tokens):
7         return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

在上述代码中,第4行便是根据指定的词表大小和维度建立一个字符嵌入层。第6~7行是字符嵌入层对应的前向传播过程,即取每个词的索引对应的词向量,其中输入tokens的形状为[len, batch_size],返回结果的形状为[len, batch_size, emb_size]

2. Positional Embedding实现

10.2.4节内容中我们已经详细介绍了位置编码的作用和原理,同时根据式(10-9)我们还能得到

$$ 1/\left(10000^{\frac{2i}{d_{\text{model}}}}\right)=\exp\left(2i\left(\frac{-\log(10000)}{d_{\text{model}}}\right)\right)\tag{10-11} $$

根据式(10-9)和式(10-11),位置编码的实现过程如下所示:

 1 class PositionalEncoding(nn.Module):
 2     def __init__(self, d_model, dropout=0.1, max_len=5000):
 3         super(PositionalEncoding, self).__init__()
 4         self.dropout = nn.Dropout(p=dropout)
 5         pe = torch.zeros(max_len, d_model) 
 6         position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  
 7         div_term = torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))
 8         pe[:, 0::2] = torch.sin(position * div_term)  
 9         pe[:, 1::2] = torch.cos(position * div_term)
10         pe = pe.unsqueeze(1)
11         self.register_buffer('pe', pe)
12     def forward(self, x): 
13         x = x + self.pe[:x.size(0), :] 
14         return self.dropout(x) 

在上述代码中,第1行d_model表示指定模型的维度,max_len表示指定最大的位置长度,需要大于最大句子的长度。第5行用于初始化一个全0的位置矩阵来保存位置信息,即图10-17中括号里的第2个矩阵。第6~9行是用来计算每个维度(每一列)的相关位置信息,其中第8行是计算偶数列的位置信息,第9行是计算奇数列的位置信息。第10行是进行维度扩充形状为[max_len, 1, d_model]。第11行是声明一个不可训练的模型变量,后续可通过以类成员变量的方式进行引用。第12~14行为前向传播计算过程,其中输入x的形状为[x_len, batch_size, emb_size],输出的形状同为[x_len, batch_size, emb_size]

3. 使用示例

在实现完这个两部分的代码之后,可以通过如下方式进行使用:

1 if __name__ == '__main__':
2     x = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], dtype=torch.long)
3     x = x.transpose(0, 1) 
4     token_embedding = TokenEmbedding(vocab_size=11, emb_size=512)
5     x = token_embedding(tokens=x)
6     pos_embedding = PositionalEncoding(d_model=512)
7     x = pos_embedding(x=x)

在上述代码中,第2~3行是指定输入,并将batch_size放到第2个维度。第4~5行是字符嵌入的计算过程,输出结果形状为[5, 2, 512]。第6~7行是位置编码的计算过程,输出结果形状同样为[5, 2, 512]

10.4.2 编码器实现#

根据图10-24可知,Transformer编码器是由多个编码层TransformerEncoderLayer所构成,因此下面我们先实现编码层然后再基于编码层构建得到编码器。

1. 编码层实现

对于一个单独的编码层来说其内部结构即为图10-18的左侧(不包括嵌入层)部分。对于这部分前向传播过程,可以通过如下代码来进行实现:

 1 class MyTransformerEncoderLayer(nn.Module):
 2     def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
 3         super(MyTransformerEncoderLayer, self).__init__()
 4         self.self_attn = MyMultiheadAttention(d_model, nhead, dropout=dropout)
 5         self.dropout1 = nn.Dropout(dropout)
 6         self.norm1 = nn.LayerNorm(d_model)
 7         self.linear1 = nn.Linear(d_model, dim_feedforward)
 8         self.dropout = nn.Dropout(dropout)
 9         self.linear2 = nn.Linear(dim_feedforward, d_model)
10         self.activation = nn.ReLU()
11         self.dropout2 = nn.Dropout(dropout)
12         self.norm2 = nn.LayerNorm(d_model)
13         
14     def forward(self, src, src_mask=None, src_key_padding_mask=None):
15         src2 = self.self_attn(src, src, src, src_mask,src_key_padding_mask)[0]
16         src = src + self.dropout1(src2)
17         src = self.norm1(src)
18         src2 = self.activation(self.linear1(src))
19         src2 = self.linear2(self.dropout(src2))
20         src = src + self.dropout2(src2)
21         src = self.norm2(src)
22         return src  

在上述代码中,第4行是实例化一个多头注意力层。第5~12行是分别实例化相关线性层、丢弃层或层归一化层。第14行中src表示经过嵌入层后的输出,形状为[src_len,batch_size, d_model]src_mask表示注意力掩码此时为Nonesrc_key_padding_mask为编码器对应的填充掩码,形状为[batch_size, src_len]。第15行是多头注意力机的前向传播计算过程,此时只取返回的第1个结果,形状为[src_len,batch_size,d_model],即多头的线性组合输出。第22行则是最后计算得到的输出,形状为[src_len,batch_size,d_model]

2. 编码器实现

在实现完一个标准的编码层后便可以基于此来实现堆叠多个编码层,从而得到Transformer中的编码器。对于这部分内容,可以通过如下代码来实现:

 1 def _get_clones(module, N):
 2     return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
 3 
 4 class MyTransformerEncoder(nn.Module):
 5     def __init__(self, encoder_layer, num_layers, norm=None):
 6         super(MyTransformerEncoder, self).__init__()
 7         self.layers = _get_clones(encoder_layer, num_layers)
 8         self.num_layers = num_layers
 9         self.norm = norm
10 
11     def forward(self, src, mask=None, src_key_padding_mask=None):
12         output = src
13         for mod in self.layers:
14             output = mod(output, mask,src_key_padding_mask)
15         if self.norm is not None:
16             output = self.norm(output)
17         return output 

在上述代码中,第1~2行是用来定义一个克隆多个编码层或解码层功能函数。第7行中的encoder_layer是一个实例化的编码层,self.layers中保存的是一个包含有多个编码层的ModuleList。第13~14行便是用来实现多个编码层堆叠起来的效果,并完成整个前向传播的计算过程。第15~17行分别是对多个编码层的输出结果进行层归一化,然后返回最后的结果,形状为[src_len,batch_size,d_model]

3. 使用示例

45