Transformer源码解读

Transformer源码解读

文章目录

• 1. 模型
• 2. 逐位前馈网络
• 3. 残差连接和层规范化
• 4. 编码器
• 5. 解码器
• 6. 训练
• 7. 小结

1. 模型

• Transformer架构
  基于编码器-解码器架构来处理序列对；跟使用注意力的seq2seq不同，Transformer是纯基于注意力的。也就是说Transformer里面没有RNN之类的；
• 基于注意力seq2seq:
  
• Transformer:
  
• 注意点：
  （1）源数据(目标数据)先进入嵌入层后和位置编码相加得到结果后再进入编码器(解码器)
  （2）编码器和解码器是可以进行N次的叠加的
  

（3）transformer的编码器是由多个相同的层叠加而成的，每个层都有两个子层，第一个子层是由多头注意力汇聚而成，第二个子层是逐位前馈网络；具体来说，在计算编码器的自注意力时，查询、键和值都是来自前一个编码器层的输出，且自注意力是queries=keys=values。每个子层受到Resnet残差网络的影响，为了让网络做得更深，编码器也引入了残差思想。对于每一个输入X,进入块后得到sublayer(X),结果满足 X+ sublayer(X),最后应用layernorm

（4）transformer的解码器也是由多个相同的层叠加起来的，并且层中也使用了残差连接和层规范化(Add+LayerNorm),解码器层由三部分组成，第一子层是遮掩的多头注意力，第二子层是多头注意力层(此层的queries来自于上一层的解码器，keys和values来自于编码器的输出)，第三子层是逐位的前馈网络

2. 逐位前馈网络

逐位前馈网络本质上就是一个MLP，也就名字起得好而已；组成 MLP -> RELU -> MLP

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: zc
# @File name: transformer_test
# @Create time: 2022/2/28 21:24
import torch
from torch import nnclass PositionalWiseFFN(nn.Module):"""基于位置的前馈网络"""def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):super(PositionalWiseFFN, self).__init__(**kwargs)# 第一个全连接层，改变输入X的最后一维度 ffn_num_input -> ffn_num_hiddens# 比如输入 X=(2,3,4),self.dense1=nn.Linear(4,8) ; X_output1=(2,3,8)self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)# relu函数lu = nn.ReLU()# 第二个全连接层，改变输入X的最后一维度 ffn_num_hiddens -> ffn_num_outputs# 比如输入 X=(2,3,8),self.dense1=nn.Linear(8,5) ; X_output1=(2,3,5)self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):# 数据流向  X -> self.dense1 -> lu -> self.dense2return self.lu(self.dense1(X)))ffn = PositionalWiseFFN(ffn_num_input=4, ffn_num_hiddens=8, ffn_num_outputs=5)
ffn.eval()
input = s((2, 3, 4))
output = ffn(input)
print(f"output.shape={output.shape}")
"""输出结果如下"""
# output.shape=torch.Size([2, 3, 5])

3. 残差连接和层规范化

层规范化(batchNormalize)是基于特征维度进规范化的。尽管批量规范化在计算机视觉中被广泛应用，但在自然语言处理任务中(输入通常是变长序列)批量规范化通常不如层规范化(LayerNormalize)的效果好

import torch
from torch import nnclass AddNorm(nn.Module):"""残差连接后进行层规范化"""def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)# 定义层规范化self.ln = nn.LayerNorm(normalized_shape)# 定义dropoutself.dropout = nn.Dropout(dropout)def forward(self, X, Y):# X为输入，Y为X经过神经网络后的输出Y=sublayer(X)# 残差连接： X + self.dropout(Y)# 层归一化为 self.ln# 流向  X + sublayer(x) -> layernorm# 残差连接要求 X.shape = Y.shapereturn self.ln((X + self.dropout(Y)))add_norm = AddNorm([3, 4], 0.5)
add_norm.eval()
input1 = s(2,3,4)
input2 = s(2,3,4)
output = add_norm(input1,input2)print(f"input1.shape={input1.shape}")
print(f"input2.shape={input2.shape}")
print(f"output.shape={output.shape}")"""输出结果如下"""
# input1.shape=torch.Size([2, 3, 4])
# input2.shape=torch.Size([2, 3, 4])
# output.shape=torch.Size([2, 3, 4])

4. 编码器

transformer的编码器包含两个子层：多头注意力和基于位置的前馈网络，这两个子层都使用了残差连接和紧随的层规范化；

• 单个编码器块EncoderBlock
  

class EncoderBlock(nn.Module):"""定义transformer的单个编码器块(EncoderBlock)"""def __init__(self, key_size, query_size, value_size, num_hiddens,num_heads, dropout, normalized_shape, ffn_num_inputs, ffn_num_hiddens,use_bias=False, **kwargs):super(EncoderBlock, self).__init__(**kwargs)# 定义多头注意力self.attention = d2l.MultiHeadAttention(key_size=key_size, query_size=query_size, value_size=value_size,num_hiddens=num_hiddens, num_heads=num_heads, dropout=dropout,bias=use_bias)# 定义第一层的加&规范化self.addnorm1 = AddNorm(normalized_shape, dropout)# 定义第二层的基于位置的前馈网络self.ffn = PositionalWiseFFN(ffn_num_inputs, ffn_num_hiddens, num_hiddens)# 定义第二层的加&规范化self.addnorm2 = AddNorm(normalized_shape, dropout)def forward(self, X, valid_lens):# 定义第一个子层：此时queries=keys=values=X来表示自注意力进入到多头注意力中在进入AddNorm层Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))# 定义第二个子层return self.addnorm2(Y, self.ffn(Y))#@save
class TransformerEncoder(d2l.Encoder):"""transformer编码器"""def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, use_bias=False, **kwargs):super(TransformerEncoder, self).__init__(**kwargs)self.num_hiddens = bedding = nn.Embedding(vocab_size, num_hiddens)self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block"+str(i),EncoderBlock(key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, use_bias))def forward(self, X, valid_lens, *args):# 因为位置编码值在-1和1之间，# 因此嵌入值乘以嵌入维度的平方根进行缩放，# 然后再与位置编码相加。X = self.pos_bedding(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

• n个编码块组成transformer的编码器

5. 解码器

import torch
from torch import nn
from d2l import torch as d2lclass AddNorm(nn.Module):"""Residual connection followed by layer normalization.Defined in :numref:`sec_transformer`"""def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)self.dropout = nn.Dropout(dropout)self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):return self.ln(self.dropout(Y) + X)class PositionWiseFFN(nn.Module):"""Positionwise feed-forward network.Defined in :numref:`sec_transformer`"""def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,**kwargs):super(PositionWiseFFN, self).__init__(**kwargs)self.dense1 = nn.Linear(ffn_num_input, ffn_num_lu = nn.ReLU()self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):return self.lu(self.dense1(X)))class DecoderBlock(nn.Module):"""transformer第 i 个解码器块代码"""def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, i, **kwargs):super(DecoderBlock, self).__init__(**kwargs)# 第 i 个解码器编号self.i = i# 掩蔽多头注意力self.attention1 = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout)# 加&规范化self.addnorm1 = AddNorm(norm_shape, dropout)# 多头注意力，query是来自上一个解码器块，key-value 来自于编码器的输出self.attention2 = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout)# 加&规范化self.addnorm2 = AddNorm(norm_shape, dropout)# 逐位前馈网络self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,num_hiddens)# 加&规范化self.addnorm3 = AddNorm(norm_shape, dropout)def forward(self, X, state):# state[0],state[1]是来存储encoder的输出enc_outputs, enc_valid_lens = state[0], state[1]# state[2]是用来存储decoder的输出，# 包含着直到当前时间步第i个块解码的输出表示# 训练阶段，输出序列的所有次元都在同一时间处理# 因此state[2][self.i]初始化为None# 预测阶段，输出序列是通过词元一个接着一个解码的# 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示if state[2][self.i] is None:key_values = Xelse:key_values = torch.cat((state[2][self.i], X), axis=1)state[2][self.i] = key_aining:batch_size, num_steps, _ = X.shape# dec_valid_lens的开头：(batch_size,num_steps),# 其中每一行是 [1,2,...,num_steps]dec_valid_lens = torch.arange(1, num_steps + 1, device=X.device).repeat(batch_size, 1)else:dec_valid_lens = None# 自注意力X2 = self.attention1(X, key_values, key_values, dec_valid_lens)Y = self.addnorm1(X, X2)# 编码器-解码器注意力# enc_outputs的开头：(batch_size,num_steps,num_hiddens)Y2 = self.attetnion2(Y, enc_outputs, enc_outputs, enc_valid_lens)Z = self.addnorm2(Y, Y2)return self.addnorm3(Z, self.ffn(Z)), stateclass TransformerDecoder(d2l.AttentionDecoder):def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, **kwargs):super(TransformerDecoder, self).__init__(**kwargs)self.num_hiddens = num_hiddensself.num_layers = bedding = nn.Embedding(vocab_size, num_hiddens)self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block"+str(i),DecoderBlock(key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, i))self.dense = nn.Linear(num_hiddens, vocab_size)def init_state(self, enc_outputs, enc_valid_lens, *args):return [enc_outputs, enc_valid_lens, [None] * self.num_layers]def forward(self, X, state):X = self.pos_bedding(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_weights# “编码器－解码器”自注意力权重self._attention_weights[1][i] = blk.attention2.attention.attention_weightsreturn self.dense(X), state@propertydef attention_weights(self):return self._attention_weights

6. 训练

import matplotlib.pyplot as plt
import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l#@save
class PositionWiseFFN(nn.Module):"""基于位置的前馈网络"""def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,**kwargs):super(PositionWiseFFN, self).__init__(**kwargs)self.dense1 = nn.Linear(ffn_num_input, ffn_num_lu = nn.ReLU()self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)def forward(self, X):return self.lu(self.dense1(X)))#@save
class AddNorm(nn.Module):"""残差连接后进行层规范化"""def __init__(self, normalized_shape, dropout, **kwargs):super(AddNorm, self).__init__(**kwargs)self.dropout = nn.Dropout(dropout)self.ln = nn.LayerNorm(normalized_shape)def forward(self, X, Y):return self.ln(self.dropout(Y) + X)#@save
class EncoderBlock(nn.Module):"""transformer编码器块"""def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, use_bias=False, **kwargs):super(EncoderBlock, self).__init__(**kwargs)self.attention = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout,use_bias)self.addnorm1 = AddNorm(norm_shape, dropout)self.ffn = PositionWiseFFN(ffn_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))#@save
class TransformerEncoder(d2l.Encoder):"""transformer编码器"""def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, use_bias=False, **kwargs):super(TransformerEncoder, self).__init__(**kwargs)self.num_hiddens = bedding = nn.Embedding(vocab_size, num_hiddens)self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block"+str(i),EncoderBlock(key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, use_bias))def forward(self, X, valid_lens, *args):# 因为位置编码值在-1和1之间，# 因此嵌入值乘以嵌入维度的平方根进行缩放，# 然后再与位置编码相加。X = self.pos_bedding(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 Xclass DecoderBlock(nn.Module):"""解码器中第i个块"""def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, i, **kwargs):super(DecoderBlock, self).__init__(**kwargs)self.i = iself.attention1 = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout)self.addnorm1 = AddNorm(norm_shape, dropout)self.attention2 = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout)self.addnorm2 = AddNorm(norm_shape, dropout)self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,num_hiddens)self.addnorm3 = AddNorm(norm_shape, dropout)def forward(self, X, state):enc_outputs, enc_valid_lens = state[0], state[1]# 训练阶段，输出序列的所有词元都在同一时间处理，# 因此state[2][self.i]初始化为None。# 预测阶段，输出序列是通过词元一个接着一个解码的，# 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示if state[2][self.i] is None:key_values = Xelse:key_values = torch.cat((state[2][self.i], X), axis=1)state[2][self.i] = key_aining:batch_size, num_steps, _ = X.shape# dec_valid_lens的开头:(batch_size,num_steps),# 其中每一行是[1,2,...,num_steps]dec_valid_lens = torch.arange(1, num_steps + 1, device=X.device).repeat(batch_size, 1)else:dec_valid_lens = None# 自注意力X2 = self.attention1(X, key_values, key_values, dec_valid_lens)Y = self.addnorm1(X, X2)# 编码器－解码器注意力。# enc_outputs的开头:(batch_size,num_steps,num_hiddens)Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)Z = self.addnorm2(Y, Y2)return self.addnorm3(Z, self.ffn(Z)), stateclass TransformerDecoder(d2l.AttentionDecoder):def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, **kwargs):super(TransformerDecoder, self).__init__(**kwargs)self.num_hiddens = num_hiddensself.num_layers = bedding = nn.Embedding(vocab_size, num_hiddens)self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)self.blks = nn.Sequential()for i in range(num_layers):self.blks.add_module("block"+str(i),DecoderBlock(key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, i))self.dense = nn.Linear(num_hiddens, vocab_size)def init_state(self, enc_outputs, enc_valid_lens, *args):return [enc_outputs, enc_valid_lens, [None] * self.num_layers]def forward(self, X, state):X = self.pos_bedding(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_weights# “编码器－解码器”自注意力权重self._attention_weights[1][i] = blk.attention2.attention.attention_weightsreturn self.dense(X), state@propertydef attention_weights(self):return self._attention_weightsnum_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, _gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)encoder = TransformerEncoder(len(src_vocab), key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,num_layers, dropout)
decoder = TransformerDecoder(len(tgt_vocab), key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)
ain_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)plt.show()

loss 0.033, 4159.7 tokens/sec on cuda:0

7. 小结

• transformer是编码器-解码器架构的一个实践，尽管在实际情况中编码器或解码器可以单独使用
• 在transformer中，多头自注意力用于表示输入序列和输出序列，不过解码器必须通过掩蔽机制来保留自回归属性