赞
踩
Transformer是当前自然语言处理乃至整个深度学习领域极其重要的经典模型,是当前大规模预训练语言模型如BERT、GPT、BART等的基本架构,读懂该模型的源码,对于深入了解其原理、后续微调、改进以该模型为基本架构的模型很有帮助,解读该代码重点在于:
另:
本文以一个简单的德语到英语的机器翻译任务Demo为例。
词嵌入本身是一个look up查表的过程,因此需要构建词表:token及其索引。现在的实际任务中,一般使用Huggingface Transformers库的Tokenizer等API直接获取。
src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4}
src_vocab_size = len(src_vocab)
tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'S': 5, 'E': 6}
tgt_vocab_size = len(tgt_vocab)
在实际任务中,应该从数据集中读取,然后构建DataLoader,本文为便于解读,仅实现一个toy过程。
‘S’(Start)表示开始字符,‘E’(End)表示结束字符,‘P’(Pad)表示填充字符。
输入文本是字符串类型的,需要将其转换成字符在词汇表中的索引,然后转换成Tensor类型。
sentences = ['ich mochte ein bier P', 'S i want a beer', 'i want a beer E']
def make_batch(sentences):
# 把文本转成词表索引
input_batch = [[src_vocab[n] for n in sentences[0].split()]]
output_batch = [[tgt_vocab[n] for n in sentences[1].split()]]
target_batch = [[tgt_vocab[n] for n in sentences[2].split()]]
# 把索引转成tensor类型
return torch.LongTensor(input_batch), torch.LongTensor(output_batch), torch.LongTensor(target_batch)
部分重要的模型超参数设置,包括:
src_len = 5 # length of source
tgt_len = 5 # length of target
## 模型参数
d_model = 512 # Embedding Size
d_ff = 2048 # FeedForward dimension
d_k = d_v = 64 # dimension of K(=Q), V
n_layers = 6 # number of Encoder of Decoder Layer
n_heads = 8 # number of heads in Multi-Head Attention
Transformer的整个网络结构包括三部分:编码层、解码层、输出层。
流程
数据形状
enc_inputs:[batch_size,src_len]
dec_inputs:[batch_size,tgt_len]
enc_outputs:[batch_size,src_len,d_model]
enc_self_attns:[batch_size,n_heads,src_len,src_len]
dec_outputs:[batch_size,tgt_len,d_model]
dec_self_attns:[batch_size,n_heads,tgt_len,tgt_len]
dec_enc_attns:[batch_size,n_heads,tgt_len,src_len]
dec_logits:[batch_size,tgt_len,tgt_vocab_size]
class Transformer(nn.Module): def __init__(self): super(Transformer, self).__init__() # 编码器 self.encoder = Encoder() # 解码器 self.decoder = Decoder() # 输出层,d_model是解码层每个token输出的维度大小,之后会做一个tgt_vocab_size大小的softmax # 因为解码输出过程相当于是在一个词表大小级别上的分类 self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False) def forward(self, enc_inputs, dec_inputs): # 输入输出部分具体要输入和返回什么参数,可以根据自己的任务和改进需要进行自定义修改,内部的执行过程是不变的 # enc_outputs就是编码器的输出,enc_self_attns是QK转置相乘之后softmax之后的注意力矩阵,代表的是每个单词和其他单词相关性; # 由于多头注意力机制会分头计算注意力,所以注意力权重矩阵是个四维向量, # 即[batch_size,n_heads,src_len,src_len] enc_outputs, enc_self_attns = self.encoder(enc_inputs) # dec_self_attns类比于enc_self_attns,是查看每个单词对decoder中输入的其余单词的相关性; # dec_enc_attns是decoder中每个单词对encoder中每个单词的相关性, # 即Cross_Attention输出的注意力权重矩阵,形状为[batch_size,n_heads,tgt_len,src_len] # 注意,这里的xxx_attns都是list类型,因为在Encoder、Decoder中将每一层的注意力权重都保存下来添加到列表中返回了 dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs) # dec_outputs映射到词表大小 dec_logits = self.projection(dec_outputs) # dec_logits : [batch_size,tgt_len, tgt_vocab_size] # 这里dec_logits进行view操作主要是为了适应后面的CrossEntropyLoss API的参数要求 return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns
model = Transformer() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) enc_inputs, dec_inputs, target_batch = make_batch(sentences) for epoch in range(10): optimizer.zero_grad() outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs) # output:[batch_size x tgt_len,tgt_vocab_size] # 就这份代码而言,这里其实可以不写.contiguous(),因为target_batch这个tensor是连续的 loss = criterion(outputs, target_batch.contiguous().view(-1)) print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss)) loss.backward() optimizer.step()
编码器由N个编码层堆叠而成。
流程
输入文本的索引tensor,经过词嵌入层得到词嵌入,然后和位置编码线性相加作为输入层的最终输出;
随后,每一层的输出最为下一层编码块的输入,在每个编码块里进行注意力计算、前馈神经网络、残差连接、层归一化等操作;
最终返回编码器最后一层的输出和每一层的注意力权重矩阵。
class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() # 这个其实就是去定义生成一个词嵌入矩阵,大小是 src_vocab_size * d_model self.src_emb = nn.Embedding(src_vocab_size, d_model) # 位置编码,这里是固定的正余弦函数,也可以使用类似词向量的nn.Embedding获得一个可以更新学习的位置编码 self.pos_emb = PositionalEncoding(d_model) # 使用ModuleList对多个encoder进行堆叠,因为后续的encoder并没有使用词向量和位置编码,所以抽离出来; self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)]) def forward(self, enc_inputs): # enc_inputs形状是:[batch_size,src_len] # 下面这个代码通过src_emb,进行索引定位,enc_outputs输出形状是[batch_size, src_len, d_model] enc_outputs = self.src_emb(enc_inputs) # 位置编码和词嵌入相加,具体实现在PositionalEncoding里,enc_outputs:[batch_size,src_len,d_model] enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1) # get_attn_pad_mask是为了得到句子中pad的位置信息,以便在计算注意力时忽略pad符号 enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) enc_self_attns = [] for layer in self.layers: # 每一层的输出作为下一层的输入,enc_outputs:[batch_size,src_len,d_model] enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask) # 把每一层得到的注意力权重矩阵添加到列表里最后返回,enc_self_attn:[batch_size,src_len,src_len] enc_self_attns.append(enc_self_attn) return enc_outputs, enc_self_attns
class EncoderLayer(nn.Module):
def __init__(self):
super(EncoderLayer, self).__init__()
self.enc_self_attn = MultiHeadAttention()
self.pos_ffn = PoswiseFeedForwardNet()
def forward(self, enc_inputs, enc_self_attn_mask):
#enc_inputs形状是[batch_size x seq_len_q x d_model],注意,最初始的QKV矩阵是等同于这个输入的
enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,V
enc_outputs = self.pos_ffn(enc_outputs) # enc_outputs: [batch_size x len_q x d_model]
return enc_outputs, attn
在后面注意力机制的部分,在计算出 Q ∗ K T Q*K^T Q∗KT除以根号之后、softmax之前得到的矩阵大小为[len_input * len_input],代表每个单词对所有(包含自己)单词的影响力。这个函数用于获取一个同等大小形状的矩阵,标记哪个位置是PAD符号,之后在计算softmax之前会把这些地方置为无穷小,避免Query去关注这些无意义的PAD符号。
注意,本函数得到的矩阵形状是[batch_size x len_q x len_k],是对K中的pad符号进行标识,并没有对Q中的做标识,因为没必要。
seq_q和seq_k不一定一致,例如,在交互注意力,q来自解码端,k来自编码端,所以告诉模型编码这边pad符号信息就可以,解码端的pad信息在交互注意力层是没有用到的。
def get_attn_pad_mask(seq_q, seq_k):
batch_size, len_q = seq_q.size()
batch_size, len_k = seq_k.size()
# eq(zero) is PAD token
pad_attn_mask = seq_k.data.eq(0).unsqueeze(1) # batch_size x 1 x len_k, one is masking
# 最终得到的应该是一个最后n列为1的矩阵,即K的最后n个token为PAD。
return pad_attn_mask.expand(batch_size, len_q, len_k) # batch_size x len_q x len_k
解码器由N个解码层堆叠而成。
解码器和编码器类似,区别在于每一层解码器有两个多头注意力机制。第一个多头注意力机制中,需要对未来词进行掩码屏蔽;第二个多头注意力机制是交叉注意力机制,需要将Encoder的输出作为K、V,将Decoder上一部分的输出作为Q。其意义就在于解码时能够利用到编码得到的上下文信息。
class DecoderLayer(nn.Module):
def __init__(self):
super(DecoderLayer, self).__init__()
self.dec_self_attn = MultiHeadAttention()
self.dec_enc_attn = MultiHeadAttention()
self.pos_ffn = PoswiseFeedForwardNet()
def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)
# 这里用dec_outputs作为Q,enc_outputs作为K和V,实现交叉注意力机制
dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)
dec_outputs = self.pos_ffn(dec_outputs)
return dec_outputs, dec_self_attn, dec_enc_attn
class Decoder(nn.Module): def __init__(self): super(Decoder, self).__init__() self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model) self.pos_emb = PositionalEncoding(d_model) self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)]) # dec_inputs : [batch_size x target_len] def forward(self, dec_inputs, enc_inputs, enc_outputs): dec_outputs = self.tgt_emb(dec_inputs) # [batch_size, tgt_len, d_model] dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1) # [batch_size, tgt_len, d_model] # 获取自注意力层pad的mask矩阵,1表示被mask dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs) # 获取look ahead的mask矩阵,即让注意力机制看不到未来的单词,获得到一个上三角为1的矩阵,1表示被mask dec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs) # 两个mask矩阵相加,大于0的为1,不大于0的为0,既屏蔽了pad的信息,也屏蔽了未来时刻的信息,为1的在之后就会被fill到无限小 # 使用gt()函数,因为可能会有在两个mask都被屏蔽的情况,1+1=2 dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequent_mask), 0) # 获取交互注意力机制中的mask矩阵,decoder的输入是q,encoder的输入是k,需要知道k里面哪些是pad符号, # 注意,q肯定也是有pad符号,但是没有必要将其屏蔽 # 这里不用再把q的未来词再mask了,因为前面已经mask过一次,输出向量对应部分的值应该都是0了 dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs) dec_self_attns, dec_enc_attns = [], [] for layer in self.layers: dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask) dec_self_attns.append(dec_self_attn) dec_enc_attns.append(dec_enc_attn) return dec_outputs, dec_self_attns, dec_enc_attns
遮蔽未来词,让当前词看不到未来词。这个函数就是用来表示Decoder的输入中哪些是未来词,显然,这个Mask矩阵应该是一个上三角矩阵。
def get_attn_subsequent_mask(seq):
"""
seq: [batch_size, tgt_len]
"""
attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
# attn_shape: [batch_size, tgt_len, tgt_len]
# np.triu()返回一个上三角矩阵,自对角线k以下元素全部置为0,k=0为主对角线
subsequence_mask = np.triu(np.ones(attn_shape), k=1) # 生成一个上三角矩阵
# 如果没转成byte,这里默认是Double(float64),占据的内存空间大,浪费,用byte就够了
subsequence_mask = torch.from_numpy(subsequence_mask).byte()
return subsequence_mask # [batch_size, tgt_len, tgt_len]
e
−
2
i
d
m
o
d
e
l
∗
l
o
g
10000
=
1
1000
0
2
i
d
m
o
d
e
l
e^{\frac{-2i}{d_{model}}*log10000}=\frac{1}{10000^{\frac{2i}{d_{model}}}}
edmodel−2i∗log10000=10000dmodel2i1
位置编码的实现直接对照着公式写就行,下面这个代码只是其中一种实现方式;
需要注意的是偶数和奇数在公式上有一个共同部分,这里使用指数函数e和log函数(以e为底)把次方拿下来,方便计算;
pos代表的是单词在句子中的绝对索引位置,例如max_len是128,那么索引就是从0,1,2,…,127,假设d_model是512,即用一个512维tensor来编码一个索引位置,那么0<=2i<512,则0<=i<=255,那么2i对应取值就是0,2,4…510,即偶数位置;2i+1的取值是1,3,5…511,即奇数位置。
最后的文本嵌入表征是词嵌入和位置编码相加得到。
class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) # 生成一个形状为[max_len,d_model]的全为0的tensor pe = torch.zeros(max_len, d_model) # position:[max_len,1],即[5000,1],这里插入一个维度是为了后面能够进行广播机制然后和div_term直接相乘 # 注意,要理解一下这里position的维度。每个pos都需要512个编码。 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # 共有项,利用指数函数e和对数函数log取下来,方便计算 div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # 这里position * div_term有广播机制,因为div_term的形状为[d_model/2],即[256],符合广播条件,广播后两个tensor经过复制,形状都会变成[5000,256],*表示两个tensor对应位置处的两个元素相乘 # 这里需要注意的是pe[:, 0::2]这个用法,就是从0开始到最后面,补长为2,其实代表的就是偶数位置赋值给pe pe[:, 0::2] = torch.sin(position * div_term) # 同理,这里是奇数位置 pe[:, 1::2] = torch.cos(position * div_term) # 上面代码获取之后得到的pe:[max_len*d_model] # 下面这个代码之后,我们得到的pe形状是:[max_len*1*d_model] pe = pe.unsqueeze(0).transpose(0, 1) # 定一个缓冲区,其实简单理解为这个参数不更新就可以,但是参数仍然作为模型的参数保存 self.register_buffer('pe', pe) def forward(self, x): """ x: [seq_len, batch_size, d_model] """ # 这里的self.pe是从缓冲区里拿的 # 切片操作,把pe第一维的前seq_len个tensor和x相加,其他维度不变 # 这里其实也有广播机制,pe:[max_len,1,d_model],第二维大小为1,会自动扩张到batch_size大小。 # 实现词嵌入和位置编码的线性相加 x = x + self.pe[:x.size(0), :] return self.dropout(x)
这里实际代码可能跟一些原理的讲解略有出入。
先映射,后分头。即拿到输入以后,映射到d_k * n_heads维度上,然后再通过转置分成n_heads个头,这样就不用写n_heads次参数矩阵,也不用进行拼接操作。
多头注意力机制完成之后还要经过一个线性层。
class MultiHeadAttention(nn.Module): def __init__(self): super(MultiHeadAttention, self).__init__() # Wq,Wk,Wv其实就是一个线性层,用来将输入映射为K、Q、V # 这里输出是d_k * n_heads,因为是先映射,后分头。 self.W_Q = nn.Linear(d_model, d_k * n_heads) self.W_K = nn.Linear(d_model, d_k * n_heads) self.W_V = nn.Linear(d_model, d_v * n_heads) self.linear = nn.Linear(n_heads * d_v, d_model) self.layer_norm = nn.LayerNorm(d_model) def forward(self, Q, K, V, attn_mask): # attn_mask:[batch_size,len_q,len_k] # 输入的数据形状: Q: [batch_size x len_q x d_model], K: [batch_size x len_k x d_model], # V: [batch_size x len_k x d_model] residual, batch_size = Q, Q.size(0) # (B, S, D) -proj-> (B, S, D) -split-> (B, S, H, W) -trans-> (B, H, S, W) # 分头;一定要注意的是q和k分头之后维度是一致的,所以一看这里都是d_k # q_s: [batch_size x n_heads x len_q x d_k] q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2) # k_s: [batch_size x n_heads x len_k x d_k] k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2) # v_s: [batch_size x n_heads x len_k x d_v] v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2) # attn_mask:[batch_size x len_q x len_k] ---> [batch_size x n_heads x len_q x len_k] # 就是把pad信息复制n份,重复到n个头上以便计算多头注意力机制 attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # 计算ScaledDotProductAttention # 得到的结果有两个:context: [batch_size x n_heads x len_q x d_v], # attn: [batch_size x n_heads x len_q x len_k] context, attn = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask) # 这里实际上在拼接n个头,把n个头的加权注意力输出拼接,然后过一个线性层,context变成 # [batch_size,len_q,n_heads*d_v]。这里context需要进行contiguous,因为transpose后源tensor变成不连续的 # 了,view操作需要连续的tensor。 context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v) output = self.linear(context) # 过残差、LN,输出output: [batch_size x len_q x d_model]和这一层的加权注意力表征向量 return self.layer_norm(output + residual), attn
这里有个关键细节就是对Mask矩阵和注意力权重矩阵的联合操作,要把注意力权重矩阵的PAD部分变成无限小,以对Query起到屏蔽作用。
class ScaledDotProductAttention(nn.Module): def __init__(self): super(ScaledDotProductAttention, self).__init__() def forward(self, Q, K, V, attn_mask): # 输入进来的维度分别是Q:[batch_size x n_heads x len_q x d_k] K:[batch_size x n_heads x len_k x d_k] V:[batch_size x n_heads x len_k x d_v] # matmul操作即矩阵相乘 # [batch_size x n_heads x len_q x d_k] matmul [batch_size x n_heads x d_k x len_k] -> [batch_size x n_heads x len_q x len_k] scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # masked_fill_(mask,value)这个函数,用value填充源向量中与mask中值为1位置相对应的元素, # 要求mask和要填充的源向量形状需一致 # 把被mask的地方置为无穷小,softmax之后会趋近于0,Q会忽视这部分的权重 scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one. attn = nn.Softmax(dim=-1)(scores) context = torch.matmul(attn, V) # context:[batch_size,n_heads,len_q,d_k] # attn:[batch_size,n_heads,len_q,len_k] return context, attn
FeedForward实际上就是一个两层的线性层,对输入进行线性转换。Position-wise意为对每个点独立做,即对序列中的每个token独立过同一个MLP,即作用在输入的最后一个维度上。
MLP有两种实现方式,一种是通过卷积的方式实现,一种是通过线性层实现。二者的区别除了原理上,还有代码细节上,例如,Conv1d要求输入必须是[batch_size,channel,length],必须是三维tensor,而Linear要求输入是[batch_size,*,d_model],可以有多个维度。
class PoswiseFeedForwardNet(nn.Module):
def __init__(self):
super(PoswiseFeedForwardNet, self).__init__()
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, inputs):
residual = inputs # inputs : [batch_size, len_q, d_model]
# Conv1d的输入为[batch, channel, length],作用于第二个维度channel,所以这里要转置
output = nn.ReLU()(self.conv1(inputs.transpose(1, 2)))
output = self.conv2(output).transpose(1, 2)
return self.layer_norm(output + residual)
class PoswiseFeedForwardNet(nn.Module):
def __init__(self):
super(PoswiseFeedForwardNet, self).__init__()
self.fc = nn.Sequential(
nn.Linear(d_model, d_ff, bias=False),
nn.ReLU(),
nn.Linear(d_ff, d_model, bias=False))
def forward(self, inputs): # inputs: [batch_size, seq_len, d_model]
residual = inputs
output = self.fc(inputs)
return nn.LayerNorm(d_model).(output + residual) # [batch_size, seq_len, d_model]
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。