Transformer分析
Pytorch中的torch.nn.modules.Transoformer.py源码分析
转载注明出处(永远鲜红の幼月)
最近想换个思路用Transformer来做实验,但对于目前比较流行的Transformer实现还较为生疏,看到pytorch中有比较简单(主要依据17年的Attention Is All You Need论文编写,没有较大的改动)的实现方法,因此想拿来作为模板分析一下来加强对Transformer的理解
Transformer类
Transformer总述:
- 该类作为最终的集成类,通过集成TransformerEncoder和Decoder进行构建。
- Encoder和Decoder分别是通过将EncoderLayer和DecoderLayer子层进行简单的堆叠形成的Encoder编码模块和Decoder解码模块。
- EncodeLayer和DecoderLayer子层
- EncoderLayer是将MultiheadAttention和FFN集成在一起的一个子层。
- DecoderLayer是将MaskMultiheadAttention,MultiheadAttention和FFN集成在一起的一个子层。
参数介绍分两部分,分别是init的类定义参数和forward的推理参数。
分别在对应的部分进行介绍
__init__()的类定义方法
类创建时先定义EncoderLayer和DecoderLayer子层,然后通过两个Layer层分别定义Encoder和Decoder,通过Encoder和Decoder两个对象建立Transformer对象。
def __init__(...):
# d_model:输出特征的数量
# dim_feadforward: feedforward神经网络的维度
# batch_first: True | False,表示输入和输出的tensor维度是(B, seq, feature) | (seq, B, feature)
# src: 输入序列编码
# ***_mask: *** 的mask
# ***_key_padding_mask: *** keys
encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedfoward, activation, ...)
encoder_norm = LayerNorm(d_model, ...)
# 通过上面两个子层定义encoder对象。
encoder = TransformerEncoder(encoder_layer, num_encoder_layers,encoder_norm)
# decoder同理...
类定义的相关参数含义:
- d_model: encoder/decoder输入的特征维度(default=512)
- hhead: 多头注意力的多头数量(default=8)
- num_encoder_layers(num_decoder_layers): 编解码器模块所包含的子层的数量(default=6)
- dim_feedforward: feedforward模型的维度(default=2048)
- dropout: Dropout权重(default=0.1)
- activation: 编解码器中间层的激活函数,可以是relu或者gelu(default=relu)
- custom_encoder(custom_decoder): (default=None)
- layer_norm_eps: LayerNorm的eps value(default=1e-5)
- batch_first: batch维度的位置,True的时候是(B,seq,feature)(default=False(seq,B,feature))
- norm_first: 使用Post-LN还是Pre-LN(区别在下文会介绍)
forward结构
def forward(src, tgt, src_mask, tgt_mask, ...):
# 前面几个RuntimeError的判定,就是一些输入参数格式不对的判定
memory = encoder(src, src_mask,...)
output = decoder(tgt, memory, tgt_mask...)
return output
简单的通过encoder和decoder的forward函数运算,然后获得输出即可。
forward推理相关参数含义:
- src: 编码器的输入序列(无batch的tensor为(S,E),有batch的tensor默认为(S,N,E))
- tgt: 解码器的输入序列(无batch的tensor为(T,E),有batch的tensor默认为(T,N,E))
- src_mask: src的附加mask矩阵(tensor为(S,S)或($ N \cdot \text{num_heads} $),S,S))
- tgt_mask: tgt的附加mask矩阵(tensor为(T,T)或($ N \cdot \text{num_heads} $),T,T))
- memory_mask: 编码器的输出mask矩阵(tensor为(T,S))
- src_key_padding_mask: src key per batch的ByteTensor mask(无batch的tensor为(S),有batch的tensor为(N,S))
- tgt_key_padding_mask: tgt key per batch的ByteTensor mask(无batch的tensor为(T),有batch的tensor为(N,T))
- memory_key_padding_mask: memory key per batch的ByteTensor mask(无batch的tensor为(S),有batch的tensor为(N,S))
例
src = torch.rand((10, 32, 512))
tgt = torch.rand((20, 32, 512))
transformer_model = nn.Transformer(nhead=16, num_encoder_layers=12)
output = tansformer_model(src, tgt)
Transformer Encoder类
该类通过encoder_layer,norm和num_layers来堆叠构建n个EncoderLayer子层来构建Encoder这一模块
__init__()的类定义方法
通过接受传入的encoder_layer和norm进行定义
def __init_(encoder_layer, num_layers):
# 通过_get_clones()函数进行简单的子层堆叠定义
self.layers = _get_clones(encoder_layer, num_layers)
# 剩下的就是对参数进行self.参数 进行私有化
...
## 补充_get_clones()函数,这个函数定义在最外面,算是一个全局函数吧,下面的decoder也调用了一次。
## 这个函数实现的就是将N个module进行堆叠,然后返回一个N个module的列表
def _get_clones(module, N):
return ModuleList([copy.deepcopy(module) for i in range(N)])
encoder的类定义方法比较简单,就是将EncoderLayer进行堆叠在一起,构成一个大的Encoder编码器就好了,主要的计算还是发生在EncoderLayer中的
forward结构
Encoder中forward看起来挺唬人,其实就是加了个特殊条件的优化判断,当触发特殊条件时,将会使用NestedTensor进行加速。
def forward(src, mask, ...):
output = src
# 这里是进行的特殊条件判断,满足特殊条件之后,会使用NestedTensor变量进行加速,详情看下文的EncoderLayer部分解释
first_layer = self.layers[0]
... # 一些特殊判定
# 最后使用
# torch._nested_tensor_from_mask(output, src_key_padding_mask.logical_not())生成NestedTensor送入EncoderLayer中进行加速
# 下面就是常规的逐层迭代运算了
for mod in self.layers:
if convert_to_nested: # 对于特殊判定的目标,传入不同的参数
output = mod(output, src_mask)
else:
output = mod(output, src_mask, ...)
if convert_to_nested:
output = output.to_padded_tensor(0.)
if self.norm is not None:
output = self.norm(output)
return output
对于前面的一堆if判定,是满足以下条件时,forward会通过特殊的优化方法实现。
- autograd被禁用(使用torch.inference_mode 或 torch.no_grad 或者参数 requies_grad没有用)
- training被禁用(用.eval()功能)
- batch_first为True并且输入是batch,(如src.dim() == 3)
- norm_first为False
- activation是relu,gelu,torch.functional.relu 或 torch.functional.gelu之一的情况
- src_mask和src_key_padding_mask参数只被传了一个
- src是NestedTensor类型,并且src_mask和src_key_padding_mask没有被传
- 两个LN有相同的eps值(默认是相同的,除非手动调节)
如果使用优化实现,则可以把src类型变为NestedTensor,表示padding比padding mask效率更高,这种情况下,将返回一个NestedTensor,可以产生一个和padding的输入部分成比例的额外加速。
具体细节可以看文档https://pytorch.org/docs/stable/nested.html
这里可以提前瞄一眼下面的EncoderLayer类,仔细对照一下就会发现,Encoder中的if判断条件和EncoderLayer中的if判断条件基本一样,把Encoder中的first_layer改成self,就跟EncoderLayer中的一样了。因为这里first_layer其实就是EncoderLayer,所以判断的条件是一样的。
TransformerDecoder类
通过n个DecoderLayer构成一个decoder模块
__init__()的类定义方法
这里主要就是按照规定的层数将decoder layer子层堆叠起来形成一个decoder解码器
def __init__(...):
self.layers = _get_clones(deocder_layer, num_layers)
# 下面就是几个参数的私有化,没什么意思
...
forward结构
没什么花活,就是遍历一下每一层,然后迭代传参计算输出即可
def forward(tgt, memory, tgt_mask, memory_mask, ...):
output = tgt
for mod in self.layers: # 迭代传参计算
output = mode(output, memory, tgt_mask, memory_mask, ...)
if self.norm is not None: # 归一化
output = self.norm(output)
return output
例
decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
transformer_decoder = nn.TransformerDeocder(decoder_layer, num_layers=6)
memory = torch.rand(10, 32, 512)
tgt = torch.rand(20, 32, 512)
out = transformer_decoder(tgt, memory)
TransformerEncoderLayer类
该类主要是对self-attn和feedforward进行了集成,然后在forward部分增加了一个对Post-LN和Pre-LN的使用区别。
另外就是补充了一下满足特殊条件时的优化加速方法。
__init__()的类定义方法
def __init__(...):
# 多头注意力
self.self_attn = MultiheadAttention(d_model, nhead, dropout, batch_first, ...)
# feedforward部分
self.linear1 = Linear(d_model, dim_feadforward, ...)
self.linear2 = Linear(dim_feadforward, d_model, ...)
# 下面的就是常规的LayerNorm和Dropout部分,省略
可以看出,Layer部分主要就是将多头注意力和Feedforward两部分进行一下集成。
forward结构
这里的forward作者在实现时参考了一篇文章https://arxiv.org/pdf/2002.04745v1.pdf
简单介绍一下这篇文章:这篇文章比较了Post-Layer norm和 Pre-Layer norm的区别,证明了Layer norm放在residual block之间输出层附近参数的期望梯度较大。在较大的梯度上使用较大的学习率会使训练不稳定,因此需要热训练学习率。将LN放在residual block之间,则会让梯度表现的更好,可以取消热身训练,以及更少的超参调整和训练时间。
Post-LN和Pre-LN如图所示(图片来自论文https://arxiv.org/pdf/2002.04745v1.pdf)
所以pytorch在定义Transformer的时候考虑到了这一情况,设置了norm_first这一参数,来控制Post-LN和Pre-LN。
剩下的就是简单的MultiheadAttention+FFN使用残差块连接的操作了。
def forward(src, src_mask, ...):
# 前面的一堆if判断先不管,这是对使用优化加速的判定
if ... :
...
# 下面是正常运行时的逻辑,分别对应的Post-LN和Pre-LN实现
if self.norm_first: # Pre-LN
x = x + self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)
x = x + self._ff_block(self.norm2(x))
else: # Post-LN
x = self.norm1(x + self._sa_block(x, src_mask, src_key_padding_mask))
x = self.norm2(x + self._ff_block(x))
return x
# 下面紧接着就是_sa_block()和_saff_block()的定义,作为self-attention和ffn各自集成的部分。
def _sa_block(x, attn_mask, ...):
x = self.self_attn(x, x, x, attn_mask, ...)[0]
return self.dropout1(x)
def _ff_block(x): # FFN就是 linear + act + dropout + linear + dropout
x = self.linear2(self.dropout(self.activation(self.linear1(x))))
return self.dropout2(x)
例
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
src = torch.rand(10, 32, 512) # 当定义时参数batch_first=True时,tensor应为(32, 10, 512)
out = encoder_layer(src)
TransformerDecoderLayer
__init__()类定义方法
定义了两个MuliheadAttention类,和2个Linear用于组成FFN。
def __init__(d_model, nhead, dim_feedforward, ...):
self.self_attn = MultiheadAttention(d_model, nhead, ...)
self.multihead_attn = MulitheadAttention(d_model, nhead, ...)
self.linear1 = Linear(d_model, dim_feedforward, ...)
self.linear2 = Linear(dim_feedforward, d_model, ...)
# ... 下面就是一些Dropout和LN等定义
forward()结构
这里没什么花活,只用了一个Post-LN和Pre-LN,其他就是用的常规的mask MultiheadAttn + MultiheadAttn + FFN套路
def forward(tgt, memory, tgt_mask, memory_mask, ...):
x = tgt
if self.norm_first: # Pre-LN
x = x + self._sa_block(self.norm1(x), tgt_mask, tgt_key_padding_mask)
x = x + self._mha_block(self.norm2(x), memory, memory_mask, memory_key_padding_mask)
x = x + self._ff_block(self.norm3(x))
else: # Post-LN
x = self.norm1(x + self._sa_block(x, tgt_mask, tgt_key_padding_mask))
x = self.norm2(x + self._mha_block(x, memory, memory_mask, memory_key_padding_mask))
x = self.norm3(x + self._ff_block(x))
# 下面就是关于_sa_block, _mha_block和_ff_block的实现
def _sa_block(x, attn_mask, key_padding_mask): # self-attn
x = self.self_attn(x, x, x, attn_mask, ...)[0]
return self.dropout1(x)
def _mha_block(x, mem, attn_mask, ...): # Multihead Attn
x = self.multihead_attn(x, mem, mem, attn_mask, ...)[0]
return dropout2(x)
def _ff_block(x): # FFN
x = self.linear2(self.dropout(self.activation(self.linear1(x))))
return self.dropout3(x)
例
decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8) # 当batch_first=True时
memory = torch.rand(10, 32, 512) # tensor改为(32, 10, 512)
tgt = torch.rand(20, 32, 512) # tensor改为(32, 20, 512)
out = decoder_layer(tgt, memory)
至此,该部分(Transformer.py)的源码已经分析完了,其中涉及到的其他部分的内容,不在本文的范围内(要不然就显得文不对题),感兴趣的可以自行查阅