Transformer源码分析（torch.nn.modules.Transformer.py）

永远鲜红の幼月

已于 2022-07-07 12:25:23 修改

阅读量2.4k

点赞数 2

分类专栏：机器学习入门科普文章标签： transformer 深度学习 python

于 2022-07-06 21:27:48 首次发布

本文链接：https://blog.csdn.net/qq_40482358/article/details/125648011

版权

机器学习入门科普专栏收录该内容

11 篇文章 0 订阅

订阅专栏

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）的源码已经分析完了，其中涉及到的其他部分的内容，不在本文的范围内（要不然就显得文不对题），感兴趣的可以自行查阅

永远鲜红の幼月

关注

2
点赞
踩
14

收藏

觉得还不错? 一键收藏
0
评论
Transformer源码分析（torch.nn.modules.Transformer.py）

Pytorch中的torch.nn.modules.Transoformer.py源码分析最近想换个思路用Transformer来做实验，但对于目前比较流行的Transformer实现还较为生疏，看到pytorch中有比较简单（主要依据17年的Attention Is All You Need论文编写，没有较大的改动）的实现方法，因此想拿来作为模板分析一下来加强对Transformer的理解......
复制链接

扫一扫