入门pytorch-Transformer

前言

虽然Transformer是2017年由Google推出，如果按照读论文只读近两年的思路看，那它无疑是过时的，但可惜的是，目前很多论文的核心依然是Transformer，或者由其进行改进的，故本文使用pytorch来搭建一下Transformer这个模型

全局分析

首先，我们要从整个模型架构入手，从大的层面看这块内容，然后再开始编写代码。欧克，这里默认大家掌握了一些基础知识，Transformer是由Google于2017年的Attention Is All You Need论文上所提出来的。如下图，则是论文中所提出的整个框架，可以很清晰的看出具体所使用的组件

总的来看，两个模块，编码器、解码器，特别有自编码器的思想(换句话说，Transformer借鉴了seq2sqe，而seq2sqe天然就是自编码器的思想)，下面引用了论文当中的话.

1. Encoder: 编码器由 $N=6$ 个相同层的堆叠组成。每一层有两个子层。第一个是多头自注意机制，第二个是一个简单的，位置全连接前馈网络。我们在两个子层中的每一个周围使用残差连接，随后进行层归一化。也就是说，每个子层的输出是$LayerNorm(x+Sublayer(x))$，其中 $Sublayer(x)$ 是由子层本身实现的函数。为了促进这些剩余连接，模型中的所有子层以及嵌入层产生维度 $d_{model}=512$ 的输出。

2. Decoder: 解码器也由 $N=6$ 个相同层的堆栈组成。除了每个编码器层中的两个子层之外，解码器还插入第三子层，该第三子层对编码器堆栈的输出执行多头注意。与编码器类似，我们在每个子层周围使用残差连接，然后进行层归一化。我们还修改了解码器堆栈中的自注意子层，以防止位置注意到后续位置。该掩蔽与输出嵌入偏移一个位置的事实相结合，确保了位置i的预测可以仅依赖于小于 $i$ 的位置处的已知输出。

好了，我们了解了整个大的框架，开始接触小的组件，由上图，我们可以进行拆分，其整个框架是由嵌入层、位置编码层、（掩码）多头注意力层、前馈神经网络层、残差连接+归一化层；接下来我们分别实现这些神经网络块（层），并将其堆叠在一起，其实就是Transformer了。

好了，接下来我们就开始编写代码吧

位置编码层

下面的就是位置编码。在这里打个比方，就比如机器翻译问题上，我们需要将一句话翻译成另一句话，即A --> B，那么A和B必然是长度随机的编码，这里输入的就是（想要翻译几句话，几个单词，每个单词的特征编码）

$P{E}_{(pos,2i)}=sin(pos/10000^{2i/d_{model}})$
$P{E}_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}})$

# 位置编码
class PositionalEncoding(nn.Module):'''num_hiddens: 神经元数量（嵌入维度数量）dropout: 神经元的丢弃概率max_len: 序列最大长度'''def __init__(self, num_hiddens, dropout, max_len=1000):super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(dropout)self.P = torch.zeros((1, max_len, num_hiddens))# 等差数列X = (torch.arange(max_len, dtype=torch.float32).reshape(-1, 1)/ torch.pow(10000, torch.arange(0, num_hiddens, 2, dtype=torch.float32) / num_hiddens))self.P[:, :, 0::2] = torch.sin(X)self.P[:, :, 1::2] = torch.cos(X)def forward(self, X): # X = (批次，序列长度，特征)X = X + self.P[:, :X.shape[1], :].to(X.device)return self.dropout(X)

残差连接+归一化层

欧克，接下来我们先来看一下简单理解的层。残差连接可以理解为防止梯度消失问题，以及加快收敛的一个有效的方法。而归一化就是防止梯度爆炸的问题

# 残差连接 + 归一化层
class AddNorm(nn.Module):'''normalized_shape: 形状大小dropout: 神经元的丢弃概率'''def __init__(self, normalized_shape, dropout):super(AddNorm, self).__init__()self.drop = nn.Dropout(dropout) # 丢弃层self.layer = nn.LayerNorm(normalized_shape) # 对输入的小批量应用层规范化def forward(self, X, Y):return self.layer(self.drop(Y) + X)

前馈神经网络层

这块也比较简单，就是两层的感知机而已。

$FFN(x)=max(0,x{W}_1+b_1)W_2+b_2$

# 前馈神经网络
class PostionWiseFFN(nn.Module):'''num_input: 输入形状num_hiddens: 隐藏层形状num_ouput: 输出层形状'''def __init__(self, num_input, num_hiddens, num_ouput):super(PostionWiseFFN, self).__init__()self.liner1 = nn.Linear(num_input, num_hiddens) # 线性层1self.relu = nn.ReLU() # 激活函数self.liner2 = nn.Linear(num_hiddens, num_ouput) # 线性层2def forward(self, X):return self.liner2(self.relu(self.liner1(X)))

点积注意力机制

这里给出点积注意力机制的公式

$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d}})V$

其中Q为查询矩阵，K键矩阵，V值矩阵，它们的维度是$(n,d_k)$，$(m,d_k)$，$(m,d_v)$，m是键值对的数量，$d_k$是键矩阵和查询矩阵的特征维度，$d_v$是值矩阵的特征维度。而除于$\sqrt{d}$是为了避免梯度消失问题

class DotProductAttention(nn.Module):'''dropout: 神经元的丢弃概率'''def __init__(self, dropout):super(DotProductAttention, self).__init__()self.dropout = nn.Dropout(dropout)def sequence_mask(self, X, valid_lens, value=0): # 根据valid_lens生成掩码'''X: (批量, 最大序列长度, 特征)valid_lens: 有效长度value: 填充数据'''maxlen = X.size(1)mask = torch.arange((maxlen), dtype=torch.float32,device=X.device)[None, :] < valid_lens[:, None]X[~mask] = valuereturn Xdef masked_softmax(self, X, valid_lens):if valid_lens is None:return nn.functional.softmax(X, dim=-1)else:shape = X.shapeif valid_lens.dim == 1:valid_lens = valid_lens.repeat_interleave(valid_lens, shape[1])else:valid_lens = valid_lens.reshape(-1)X = self.sequence_mask(X, valid_lens)return nn.functional.softmax(X.reshape(shape), dim=-1)def forward(self, Q, K, V, valid_lens=None):d = Q.shape[-1]scores = torch.bmm(Q, K.transpose(1, 2)) / math.sqrt(d) # 点积获取注意力分数self.attention_weights = self.masked_softmax(scores, valid_lens) # 获取注意力权重return torch.bmm(self.dropout(self.attention_weights), V) # 注意力权重 * 值V

到这里，我们由此便可以提出自注意力机制，而自注意力机制对上面做出了一个很简单的改变，就是$Q=K=V$，换句话说Q、K、V同源

多头注意力

所谓多头注意力机制，就是有多个自注意力机制并行，然后将输出进行拼接送到线性层进一步整合

$$\begin{gathered} MultiHead(Q,K,V)=Concat(hea{d}_1,...,hea{d}_h)W^O \\ wherehea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V) \end{gathered}$$

# 多头注意力
class MultiHeadAttention(nn.Module):'''key_size: K值形状大小query_size: Q值形状大小value_size: V值形状大小num_hiddens: 隐藏层神经元数量num_heads: 头数dropout: 神经元的丢弃概率bias: 是否学习加性偏差，默认不学习'''def __init__(self, key_size, query_size, value_size, num_hiddens,num_heads, dropout, bias=False):super(MultiHeadAttention, self).__init__()assert num_hiddens % num_heads == 0self.num_heads = num_headsself.attention = DotProductAttention(dropout) # 这里使用到了点积注意力机制self.W_q = nn.Linear(query_size, num_hiddens, bias=bias) # Q矩阵权重self.W_k = nn.Linear(key_size, num_hiddens, bias=bias) # K矩阵权重self.W_v = nn.Linear(value_size, num_hiddens, bias=bias) # V矩阵权重self.W_o = nn.Linear(num_heads, num_hiddens, bias=bias) # 输出线性层def forward(self, Q, K, V, valid_lens): # valid_lens 有效长度Q = transpose_qkv(self.W_q(Q), self.num_heads)K = transpose_qkv(self.W_k(K), self.num_heads)V = transpose_qkv(self.W_v(V), self.num_heads)if valid_lens is not None:valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)output = self.attention(Q, K, V, valid_lens)output_concat = transpose_output(output, self.num_heads)return self.W_o(output_concat)

这里我们通过transpose_qkv和transpose_output函数来将数据进行改造，以此来进行并行操作

def transpose_qkv(X, num_heads):"""为了多注意力头的并行计算而变换形状"""# 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)# 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)X = X.permute(0, 2, 1, 3)return X.reshape(-1, X.shape[2], X.shape[3])
def transpose_output(X, num_heads):"""逆转transpose_qkv函数的操作"""X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])X = X.permute(0, 2, 1, 3)return X.reshape(X.shape[0], X.shape[1], -1)

编码器

如下则是编码器的代码

class EncoderBlock(nn.Module):'''K_size, Q_size, V_size: 键值对和查询的大小num_hiddens: 中间隐藏层神经元num_heads: 头数norm_shape: 残差连接+归一化层的输入形状num_input: FFN的输入形状ffn_num_hiddens: FFN隐藏层神经元dropout: 神经元的丢弃概率'''def __init__(self, K_size, Q_size, V_size, num_hiddens, num_heads,norm_shape, num_input, ffn_num_hiddens, dropout, bias=False):super(EncoderBlock, self).__init__()self.attention = MultiHeadAttention(K_size, Q_size, V_size, num_hiddens, num_heads, dropout, bias) # 多头注意力机制 self.addnorm1 = AddNorm(norm_shape, dropout) # 归一化+残差连接self.ffn = PostionWiseFFN(num_input, ffn_num_hiddens, num_hiddens) # 前馈神经网络self.addnorm2 = AddNorm(norm_shape, dropout) # 归一化+残差连接def forward(self, X, valid_lens):Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))return self.addnorm2(Y, self.ffn(Y))

接下来我们需要将位置编码层、嵌入层、残差连接+归一化层、前馈神经网络、编码器等封装起来

class TransformerEncoder(nn.Module):'''vocab_size: 词典大小K_size, Q_size, V_size: 键值对和查询的大小num_hiddens: 中间隐藏层神经元num_heads: 头数norm_shape: 残差连接+归一化层的输入形状num_input: FFN的输入形状ffn_num_hiddens: FFN隐藏层神经元num_layers: 编码器的层数dropout: 神经元的丢弃概率'''def __init__(self, vocab_size, K_size, Q_size, V_size, num_hiddens, num_heads,norm_shape, num_input, ffn_num_hiddens, num_layers, dropout, bias=False):super(TransformerEncoder, self).__init__()self.num_hiddens = num_hiddens#  num_embeddings: 嵌入词典的大小  embedding_dim: 每个嵌入向量的尺寸self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=num_hiddens) # 嵌入层self.pos_encoding = PositionalEncoding(num_hiddens, dropout) # 位置编码层self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block" + str(i),EncoderBlock(K_size, Q_size, V_size, num_hiddens, num_heads, norm_shape, num_input, ffn_num_hiddens, dropout, bias))def forward(self, X, valid_lens):# 输入 X (句子个数，单词个数)# 输出 ret (句子个数，单词个数，单词特征)X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))self.attention_weights = [None] * len(self.blks)for i, blk in enumerate(self.blks):X = blk(X, valid_lens)self.attention_weights[i] = blk.attention.attention.attention_weightsreturn X

解码器

class DecodeBlock(nn.Module):'''K_size, Q_size, V_size: 键值对和查询的大小num_hiddens: 中间隐藏层神经元num_heads: 头数norm_shape: 残差连接+归一化层的输入形状num_input: FFN的输入形状ffn_num_hiddens: FFN隐藏层神经元dropout: 神经元的丢弃概率i:'''def __init__(self, K_size, Q_size, V_size, num_hiddens, num_heads,norm_shape, num_input, ffn_num_hiddens, dropout, i):super(DecodeBlock, self).__init__()self.i = iself.attention1 = MultiHeadAttention(K_size, Q_size, V_size, num_hiddens, num_heads, dropout) # 多头注意力机制self.addnorm1 = AddNorm(norm_shape, dropout)self.attention2 = MultiHeadAttention(K_size, Q_size, V_size, num_hiddens, num_heads, dropout) # 多头注意力机制self.addnorm2 = AddNorm(norm_shape, dropout)self.ffn = PostionWiseFFN(num_input, ffn_num_hiddens, num_hiddens)self.addnorm3 = AddNorm(norm_shape, dropout)def forward(self, X, state):# state [编码器输入，编码器有效长度，中间状态用于记录]enc_outputs, enc_valid_lens = state[0], state[1]if state[2][self.i] is None:K = Xelse:K = torch.cat((state[2][self.i], X), axis=1)state[2][self.i] = Kif self.training:batch_size, num_steps, _ = X.shapedec_valid_lens = torch.arange(1, batch_size + 1, device=X.device)else:dec_valid_lens = NoneX2 = self.attention1(X, K, K, dec_valid_lens)Y = self.addnorm1(X, X2)Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)Z = self.addnorm2(Y, Y2)return self.addnorm3(Z, self.ffn(Z)), state

封装

class TransformerDecoder(nn.Module):'''vocab_size: 词典大小K_size, Q_size, V_size: 键值对和查询的大小num_hiddens: 中间隐藏层神经元num_heads: 头数norm_shape: 残差连接+归一化层的输入形状num_input: FFN的输入形状ffn_num_hiddens: FFN隐藏层神经元num_layers: 编码器的层数dropout: 神经元的丢弃概率'''def __init__(self, vocab_size, K_size, Q_size, V_size, num_hiddens, num_heads,norm_shape, num_input, ffn_num_hiddens, num_layers, dropout):super(TransformerDecoder, self).__init__()self.num_hiddens = num_hiddensself.num_layers = num_layersself.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=num_hiddens)self.pos_encoding = PositionalEncoding(num_hiddens, dropout)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block" + str(i),DecodeBlock(K_size, Q_size, V_size, num_hiddens, num_heads, norm_shape, num_input, ffn_num_hiddens, dropout, i))self.linear = nn.Linear(num_hiddens, vocab_size)def init_state(self, enc_outputs, enc_valid_lens):return [enc_outputs, enc_valid_lens, [None] * self.num_layers]def forward(self, X, state):X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))self._attention_weights = [[None] * len(self.blks) for _ in range(2)]for i, blk in enumerate(self.blks):X, state = blk(X, state)self._attention_weights[0][i] = blk.attention1.attention.attention_weightsself._attention_weights[1][i] = blk.attention2.attention.attention_weightsreturn self.linear(X), self._attention_weights@propertydef attention_weights(self):return self._attention_weights

Transformer

最后整合一下

class EncoderDecoder(nn.Module):"""编码器-解码器架构的基类"""def __init__(self, encoder, decoder, **kwargs):super(EncoderDecoder, self).__init__(**kwargs)self.encoder = encoderself.decoder = decoderdef forward(self, enc_X, dec_X, *args):enc_outputs = self.encoder(enc_X, *args)dec_state = self.decoder.init_state(enc_outputs, *args)return self.decoder(dec_X, dec_state)

Transformer以三种不同的方式使用多头注意力：(引自论文)

• 在 “encoder-decoder attention” 层中，查询$Q$来自先前的解码器层，并且存储器键$K$和值$V$来自编码器的输出。这使得解码器中的每个位置都能处理输入序列中的所有位置。这模拟了序列到序列模型中的典型编码器-解码器注意机制。
• 编码器包含自我注意层。在自关注层中，所有的键、值和查询都来自同一个地方，在这种情况下，是编码器中前一层的输出。编码器中的每个位置都可以处理编码器的前一层中的所有位置。
• 类似地，解码器中的自关注层允许解码器中的每个位置关注解码器中的直到并且包括该位置的所有位置。我们需要防止解码器中的冗余信息流，以保持自回归特性。我们通过屏蔽（设置为 $-\infty$）$softmax$ 输入中对应于非法连接的所有值来实现这一点。

最后，大家再来回顾一下Transformer模型整体架构