Transformer实现
一. PyTorch中的 Transformer 层
在 PyTorch 中,框架并未直接封装一个端到端的全自动Transformer,而是提供了各类神经网络层模块,供开发者灵活搭建。
torch.nn 模块下提供了专门服务于 Transformer 架构的各类网络层,用于构建复杂的编码器与解码器。
1. 核心组件概述
| 类名称 | 作用 |
|---|---|
nn.Transformer | 不带输入与输出层的 Transformer 模型,同时具备编码器和解码器 |
nn.TransformerEncoder | Transformer 编码器的堆叠层,可以控制 Nx 的N的具体数字 |
nn.TransformerDecoder | Transformer 解码器的堆叠层,可以控制 Nx 的N的具体数字 |
nn.TransformerEncoderLayer | Transformer 编码器单层,由自注意力和前馈网络组成 |
nn.TransformerDecoderLayer | Transformer 解码器单层,由自注意力、编码器-解码器注意力和前馈网络组成 |
nn.MultiheadAttention | 多头注意力机制模块 |
nn.LayerNorm | 层归一化(Layer Normalization)层 |
nn.Embedding | 嵌入层,用于将输入词索引转换为稠密的向量表示 |
(注:在所有这些类中,最核心和常用的是 nn.TransformerEncoderLayer 与 nn.TransformerDecoderLayer,它们赋予了架构极高的灵活性。)
2. nn.Transformer类 - 标准Transformer模型
nn.Transformer封装了完整的Transformer结构。如下图所示,它对Encoder和Decoder两部分的包装,它并没有实现输入中的Embedding和Positional Encoding和最后输出的Linear+softmax部分。
nn.Transformer类同样是nn.Module类的子类
(1).nn.Transformer类类型
作用:
nn.Transformer 是PyTorch中实现了经典论文 “Attention is All You Need” 中完整 Encoder-Decoder(编码器-解码器)架构的顶级类。它将多头自注意力机制(Multi-Head Attention)和前馈神经网络(FFN)高度封装,专门用于处理序列到序列(Seq2Seq)的任务,例如机器翻译、文本摘要等。
torch.nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
dim_feedforward=2048, dropout=0.1, activation='relu', custom_encoder=None,
custom_decoder=None, layer_norm_eps=1e-05, batch_first=False, norm_first=False,
bias=True, device=None, dtype=None)
核心参数:
- 词向量维度 (d_model): 输入序列中每个 token(词)的特征向量维度。默认是 512。(
int) - 注意力头数 (nhead): 多头注意力机制(Multi-Head Attention)中的头数。注意:
d_model必须能被nhead整除。(int) - 编码器/解码器层数 (num_encoder_layers/num_decoder_layers):Encoder和Decoder中分别堆叠的Block数量。默认为6层。(
int) - 前馈网络维度 (dim_feedforward): 每个 Transformer Block 内部两层全连接层中间的隐藏层(升维)维度。默认 2048。(
int) - 丢弃率 (dropout):Dropout出现在自注意力层后、残差链接之前,也出现在前馈神经网络后、残差链接之前。默认 0.1。(
float) - 激活函数 (activation): 内部FFN层使用的非线性激活函数,支持
'relu'或'gelu'。(str或Callable) - 批次优先开关 (batch_first): 极其关键。如果设为
False(默认),输入张量形状必须是(seq_len, batch_size, d_model);若设为True,则符合直觉的(batch_size, seq_len, d_model)。(bool) - 归一化顺序 (norm_first): 决定 LayerNorm 是放在 Attention/FFN 之前 (Pre-Norm) 还是之后 (Post-Norm)。默认
False(Post-Norm)。训练极深模型时强烈建议设为True。(bool)
不重要的参数:
- 自定义编码器 (custom_encoder): 默认情况下,PyTorch 会自动帮你堆叠官方的
TransformerEncoder。如果你设计了一种魔改版的 Encoder,可以直接通过这个参数把你的自定义模块传进去替换掉官方实现。默认为None。(nn.Module或None) - 自定义解码器 (custom_decoder): 同上,用于替换官方默认的 Decoder 模块。默认为
None。(nn.Module或None) - 层归一化微小值 (layer_norm_eps): 传入 LayerNorm 中防止除以 0 的极小值 ϵ\epsilonϵ (Epsilon),默认 1e-05。(
float) - 偏置项开关 (bias): 决定内部所有线性层和归一化层是否学习偏置项 (Bias)。默认
True。(bool) - 分配设备 (device): 指定模型初始化的权重存放的硬件设备(如
'cuda:0')。(torch.device或str或None) - 数据类型 (dtype): 指定模型权重的数据类型(如
torch.float16、torch.bfloat16)。(torch.dtype或None)
理论与底层架构的对齐 (极其重要):
在宏观架构上,nn.Transformer 是一个封装好的黑盒,它内部自动帮你完成了我们之前拆解过的两步走战略:
-
Encoder 提取特征: 接收源序列输入(源信息),并输出高维上下文表示(即Encoder输出)。
Eout=Encoder(Xsrc)E_{\text{out}} = \text{Encoder}(X_{\text{src}})Eout=Encoder(Xsrc)
-
Decoder 交叉生成: 接收目标序列输入(已生成信息)和Encoder的输出,通过Cross-Attention进行查询,最终输出预测特征。
Dout=Decoder(Xtgt,Eout)D_{\text{out}} = \text{Decoder}(X_{\text{tgt}}, E_{\text{out}})Dout=Decoder(Xtgt,Eout)
补充 (史诗级避坑警告):
Transformer 的两大夺命陷阱 (Shape & Cheat Trap):
形状隐式翻转陷阱 (
batch_first):
PyTorch 早期为了兼容 RNN 的底层内存连续性,默认所有的 Transformer 组件都是batch_first=False。必须手动加前馈掩码:
自回归 Decoder需要手动添加。
(2). _call_ - 实例调用(执行前向传播)
作用:接收源序列、目标序列以及各种掩码(Mask),执行完整的前向传播,计算出最终的输出特征图,并构建计算图。同理,作为 nn.Module 的子类,我们严禁直接调用 .forward(src, tgt),必须直接把实例对象当作函数调用。
我们常定义nn.Transformer类型对象叫model或transformer
# 实例化后,直接像调用函数一样调用该对象
model(src, tgt, src_mask=None, tgt_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None)
参数:
- 源输入张量 (src): 即
encoder_input,喂给 Encoder 的源序列特征矩阵。(Tensor) - 目标输入张量 (tgt): 即
decoder_input,喂给 Decoder 的目前已生成的目标序列特征矩阵。(Tensor) 源掩码 (src_mask): 作用于 Encoder 源序列自注意力计算的掩码,用于人为干预源序列内部的可见性,绝大多数情况不传。形状必须是(src_seq_len, src_seq_len)。用于人为限制源序列内部特定 Token 之间的互相可见性。 (Tensor或None)- 目标掩码 (tgt_mask): 作用于 Decoder 目标序列自注意力计算的自回归掩码(通常为下三角矩阵),必须传入,用于防止模型偷看未来词。形状必须是
(tgt_seq_len, tgt_seq_len)。注意其逻辑:通常上三角区域全为负无穷(代表掩盖/阻断未来信息),主对角线及下三角为 0(代表通行)。 (Tensor或None) 记忆掩码 (memory_mask): 作用于 Decoder 交叉注意力计算的掩码,用于人为限制 Decoder 对 Encoder 输出的特定位置可见性,极少使用。形状必须是(tgt_seq_len, src_seq_len)。用于人为限制交叉注意力中特定 Token 之间的互相可见性。 (Tensor或None)- 源序列补齐掩码 (src_key_padding_mask): 作用于源序列的布尔矩阵,告诉 Encoder(算自注意力时)和 Decoder(算交叉注意力时)忽略
src中为了凑长度而填充的<PAD>字符。形状必须是(batch_size, src_seq_len)。注意其反直觉逻辑:True代表该位置是无意义的<PAD>需要被掩盖/忽略,False代表真实有效的单词需要参与计算。 (Tensor或None) - 目标序列补齐掩码 (tgt_key_padding_mask): 作用于目标序列的布尔矩阵,告诉 Decoder 忽略
tgt中的<PAD>字符。形状必须是(batch_size, tgt_seq_len)。同理反直觉逻辑:True代表掩盖/忽略无意义字符,False代表真实有效。 (Tensor或None) 记忆补齐掩码 (memory_key_padding_mask): 在完整的nn.Transformer前向传播中,通常由src_key_padding_mask自动代劳,但如果你有极特殊的控制需求,可以单独传入覆盖它,作用于交叉注意力的 PAD 屏蔽。形状必须是(batch_size, src_seq_len)。同理反直觉逻辑:True代表掩盖/忽略无意义字符,False代表真实有效。 (Tensor或None)
返回值:
- 成功: 返回经过完整 Encoder-Decoder 架构处理后的最终输出张量,形状与输入的
tgt绝对相同,包含了每个时间步的特征预测向量。(Tensor)
示例:
# 以机器翻译为例
# 1.构造数据
batch_size = 2
src_seq_len = 10 # 原文长度
tgt_seq_len = 4 # 目标/已生成文本长度(训练时知道长度,推理时自回归预测)
d_model = 512
encoder_input = torch.rand(batch_size, src_seq_len, d_model) # 即 src
decoder_input = torch.rand(batch_size, tgt_seq_len, d_model) # 即 tgt
# 2.生成给Decode用的自回归掩码 (形状必须与tgt_seq_len匹配)
tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_seq_len)
# 3.实例化完整Transformer模型 (注意:开启 batch_first=True)
transformer_model = nn.Transformer(
d_model=d_model,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6,
batch_first=True
)
# 4.执行前向传播:触发 __call__
output = transformer_model(
src=encoder_input,
tgt=decoder_input,
tgt_mask=tgt_mask
)
# Transformer 完整输出形状
print(output.shape)
# torch.Size([2, 4, 512])
3. Transformer 的五大掩码(Mask)
(1).掩码家族全解析
在工业界实践中,我们将这四大类掩码归纳如下:
tgt_mask(自回归掩码 / 二维方阵):
- 生死攸关的灵魂: 必须给!用于 Decoder 进行 Self-Attention 时,防止模型“偷看未来”。
- 形状:
(tgt_seq_len, tgt_seq_len)。通常上三角区域全为-inf(代表阻断未来信息),主对角线及下三角为0(代表通行)。
src_mask/memory_mask(交叉/源注意力掩码/二维方阵):
- 极罕见的黑科技:
src_mask限制源序列内部可见性,memory_mask限制 Decoder 对 Encoder 输出的查阅范围。 - 形状: 分别为
(src_seq_len, src_seq_len)和(tgt_seq_len, src_seq_len)。 - 工程经验: 绝大多数常规任务直接传
None即可。
src_key_padding_mask(原文端补齐过滤 - 第一道关卡):
- 节省算力的清道夫: 批量训练时,句子长短不一,Encoder 中填补的
<PAD>必须被屏蔽。 - 形状:
(batch_size, src_seq_len)的布尔型矩阵。 - 作用: 专门用于 Encoder 内部计算自身的自注意力时,让模型明确“哪些位置是真正的原文,哪些位置是凑数的垃圾”。
memory_key_padding_mask(原文端补齐过滤 - 第二道关卡):
- 翻译时的防误导: 当 Decoder 在进行交叉注意力计算时,需要从Encoder的输出(
memory)中提取原文信息。此时,它必须知道原文里哪里是<PAD>。 - 形状:
(batch_size, src_seq_len)的布尔型矩阵。 - 底层联动(关键点): 在调用顶层
nn.Transformer模型时,你只需要传入src_key_padding_mask,PyTorch 底层会自动将其作为memory_key_padding_mask传递给 Decoder。不需要手动再次定义。
tgt_key_padding_mask(输出端补齐过滤):
- 作用类似src_key_padding_mask: 屏蔽Decoder句子(答案)中为了凑数而存在的
<PAD>字符。 - 形状:
(batch_size, tgt_seq_len)的布尔型矩阵。
总结:
- Encoder 算自己看自己: 用
src_key_padding_mask。 - Decoder 算自己看自己: 用
tgt_key_padding_mask。 - Decoder 算自己看 Encoder: 用
memory_key_padding_mask(会自动从src_key_padding_mask同步)。
(2).Padding Mask的布尔反转陷阱
在 PyTorch 中,所有的
key_padding_mask的布尔值逻辑是极度“反直觉”的:
True代表 “掩盖 (Mask)” = “这里是垃圾<PAD>,请忽略它”。False代表 “保留 (Keep)” = “这里是真实单词,请计算它”。永远不要习惯性地将真实的单词设为
1并转为True,否则模型会把所有的真实输入当成瞎子屏蔽掉,导致 Loss 爆炸或输出NaN。
(3).普通掩码和前瞻掩码实现函数
①.普通掩码
需升维的填充掩码函数:
def create_padding_mask(seq, pad_token=0):
# seq: (batch_size, seq_len, embedding_dim)
# 检查填充值位置
padding_mask = (seq == pad_token).all(dim=-1) # (batch_size, seq_len)
# 增加维度以匹配注意力权重矩阵的形状
# (batch_size, num_heads, seq_len, seq_len)
padding_mask = padding_mask.unsqueeze(1).unsqueeze(3).expand(-1, -1, -1, seq.size(1))
# 将填充值部分设置为负无穷大,有效数据部分设置为0
padding_mask = padding_mask.float() * -1e9 # (batch_size, num_heads, seq_len, seq_len)
return padding_mask
注:如果是配合PyTorch中已经设置好的Transformer类来使用,则二维的掩码矩阵((seq == pad_token).all(dim=-1).float())就足够了,Transformer类会自动执行将掩码矩阵升维的过程;如果是利用更底层的机制创建的Transformer,则会需要我们手动执行上述流程来匹配掩码的结构。在实际使用时,要根据实际情况选择是否主动对掩码矩阵进行升维。
无需升维的填充掩码函数:
def create_padding_mask(seq, pad_token=0):
# seq: (batch_size, seq_len, embedding_dim)
# 检查填充值位置
padding_mask = (seq == pad_token).all(dim=-1) # (batch_size, seq_len)
padding_mask = padding_mask.float() * -1e9
return padding_mask
②.前瞻掩码
def create_look_ahead_mask(seq_len, start_seq = 1):
mask = torch.triu(torch.ones((seq_len, seq_len)),diagonal=start_seq) # triu左下方的三角矩阵,diagonal控制对角线位置
mask = mask.float() * -1e9 # 将未来的位置设置为负无穷大
return mask # (seq_len, seq_len)
前瞻掩码矩阵的结构为(seq_len, seq_len),而填充掩码矩阵的结构为(batch_size,num_heads,seq_len,seq_len)。前者可以通过广播的方式与QK.TQK.TQK.T矩阵相加,后者则必须写明4个维度的信息,这是因为前瞻掩码对所有的序列都是一样的掩码方式,但填充掩码却是在每个batch内都是不一致的,因为每个batch内的句子可能会不一致。
3. nn.TransformerEncoderLayer类 - Transformer编码器层
nn.TransformerEncoderLayer与nn.TransformerDecoderLayer: 这两个类表示Transformer单一的编码器和单一的解码器(他们代表了架构图中展示的结构,而不包括Nx的部分),又叫作分割的编码器与解码器。他们都包含了自注意力机制(self-attention)、多头注意力机制(Multi-head Attention)和前馈网络(feedforward network),以及必要的归一化和残差连接。这两个层的区别在于:
- DecoderLayer默认带有teacher forcing机制,而Encoder layer则没有这个机制。
- DecoderLayer带有的Multi-head Attention层可以用来处理编码器-解码器注意力层中的运算,但是EncoderLayer中带有的多头注意力层却没有这个机制。
nn.TransformerEncoderLayer类同样是nn.Module类的子类
(1). nn.TransformerEncoderLayer类类型
作用:
nn.TransformerEncoderLayer 是构成Transformer Encoder的基本积木(Block)。它内部封装了“多头自注意力机制和前馈神经网络 并带了残差连接和层归一化。它的核心任务是对输入的序列进行深度的特征提取和上下文信息融合。
torch.nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)
核心参数:
- 词向量维度 (d_model): 输入特征的维度,即每个 Token 的向量长度 (如 512)。(
int) - 注意力头数 (nhead): 多头注意力机制的头数。要求
d_model必须能被nhead整除。(int) - 前馈网络维度 (dim_feedforward): 内部两层全连接层中间的升维大小,默认 2048。用于提供强大的非线性拟合能力。(
int) - 丢弃率 (dropout):Dropout出现在自注意力层后、残差链接之前,也出现在前馈神经网络后、残差链接之前。默认 0.1。(
float) - 激活函数 (activation): FFN(前馈网络)内部升维后使用的非线性激活函数。支持字符串 ‘relu’(默认)或 ‘gelu’,也可以直接传入一个函数对象。(
str或Callable) - 层归一化微小值 (layer_norm_eps): 传入Layer Normalization中的一个极小值 ϵ\epsilonϵ (Epsilon),默认 1e-05。它的唯一作用是加在求方差的除法分母上,防止出现“除以 0”的数学崩溃(数值稳定性保障),绝大多数情况下永远不需要改动。(
float) - 批次优先 (batch_first): 若设为
True,则输入输出张量形状为(batch, seq, feature)。(bool) - 归一化顺序 (norm_first): 决定LayerNorm是放在Attention/FFN 之前 (Pre-Norm) 还是之后 (Post-Norm)。默认
False(Post-Norm)。在训练几十层甚至上百层的深层大模型时,工业界通常将其设为True(Pre-Norm) 以保证梯度稳定。(bool) - 偏置项开关 (bias): 决定内部的线性层 (Linear) 和层归一化 (LayerNorm) 是否学习并使用偏置项 (Bias)。默认设为 True。(
bool) - 分配设备 (device): 指定该层初始化的参数权重存放在哪个硬件设备上(如 torch.device(‘cuda:0’) 或 ‘cpu’)。(
torch.device或str或None) - 数据类型 (dtype): 指定该层参数张量的数据类型(如 torch.float32,torch.float16,torch.bfloat16)。默认 None(即使用全局默认的 float32)。(
torch.dtype或None)
理论与底层代码的对齐 (极其重要):
在经典论文的理论推导中(以默认的 Post-Norm 为例),这一层的计算分为两步:
-
自注意力与残差归一化:
x′=LayerNorm(x+MultiHeadAttention(x,x,x))x' = \text{LayerNorm}(x + \text{MultiHeadAttention}(x, x, x))x′=LayerNorm(x+MultiHeadAttention(x,x,x))(注:三个 xxx 分别作为 Query, Key, Value 传入)
-
前馈网络与残差归一化:
y=LayerNorm(x′+FFN(x′))y = \text{LayerNorm}(x' + \text{FFN}(x'))y=LayerNorm(x′+FFN(x′))
(2). _call_ - 实例调用(执行前向传播)
作用:接收源序列特征,执行单层 Encoder 的前向计算。作为 nn.Module,同样严禁直接调用 .forward(src)。
encoder_layer(src, src_mask=None, src_key_padding_mask=None)
参数:
- 输入张量 (src): 需要编码的序列张量。(
Tensor) 源注意力掩码 (src_mask): 极少在 Encoder 中使用(通常为None),除非你想人为干预某些 Token 之间的可见性。形状必须是(src_seq_len, src_seq_len)。用于人为限制源序列内部特定 Token 之间的互相可见性。 (Tensor或None)- 补齐掩码 (src_key_padding_mask): 屏蔽原文中
<PAD>字符的布尔张量。形状必须是(batch_size, src_seq_len)。注意其反直觉逻辑:True代表该位置是无意义的<PAD>需要被掩盖/忽略,False代表真实有效的单词需要参与计算。 (Tensor或None)
返回值:
- 成功: 返回经过特征提取后的张量,形状与输入的
src绝对相同。(Tensor)
补充 (史诗级避坑警告):
Pad Mask 的布尔值反转陷阱:
在传入src_key_padding_mask(用于忽略句子中凑字数的<PAD>) 时,PyTorch 的逻辑是:True代表“忽略/掩盖”,False代表“保留/参与计算”。
初学者常习惯用1表示有效数据,0表示PAD,如果直接转成布尔值扔进去,会导致模型把所有真实单词都 Mask 掉,反而去对全 0 的PAD算注意力,最终输出全NaN或极其离谱的 Loss!
示例:
batch_size = 2
src_len = 10
tgt_len = 4
d_model = 512
# 1.准备数据与 Mask
encoder_input = torch.rand(batch_size, src_len, d_model)
decoder_input = torch.rand(batch_size, tgt_len, d_model)
tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_len)
# 2.实例化单层 Encoder 和 Decoder (开启 batch_first)
enc_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=8, batch_first=True)
dec_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=8, batch_first=True)
# 3.前向计算
# 第一步:Encoder 层提取原文特征
encoder_output = enc_layer(src=encoder_input)
print("Encoder Output Shape:", encoder_output.shape)
# 返回: torch.Size([2, 10, 512])
4. nn.TransformerDecoderLayer类 - Transformer解码器层
nn.TransformerEncoderLayer与nn.TransformerDecoderLayer: 这两个类表示Transformer单一的编码器和单一的解码器(他们代表了架构图中展示的结构,而不包括Nx的部分),又叫作分割的编码器与解码器。他们都包含了自注意力机制(self-attention)、多头注意力机制(Multi-head Attention)和前馈网络(feedforward network),以及必要的归一化和残差连接。这两个层的区别在于:
- DecoderLayer默认带有teacher forcing机制,而Encoder layer则没有这个机制。
- DecoderLayer带有的Multi-head Attention层可以用来处理编码器-解码器注意力层中的运算,但是EncoderLayer中带有的多头注意力层却没有这个机制。
nn.TransformerDecoderLayer类同样是nn.Module类的子类
(1). nn.TransformerDecoderLayer类类型
作用:
nn.TransformerDecoderLayer 是构成 Transformer Decoder 的基本积木。相比 Encoder 层,它多了一个交叉注意力机制。它的任务是:一边看着自己目前已经生成的内容(通过 Masked Self-Attention),一边去查阅 Encoder提取的原文特征 (通过交叉注意力机制),最终输出对下一个词的预测特征。
torch.nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', layer_norm_eps=1e-05, batch_first=False, norm_first=False, bias=True, device=None, dtype=None)
核心参数(和EncoderLayer完全一致):
- 词向量维度 (d_model): 输入特征的维度,即每个 Token 的向量长度 (如 512)。(
int) - 注意力头数 (nhead): 多头注意力机制的头数,注:这里的nhead不能控制编码解码层的头数。要求
d_model必须能被nhead整除。(int) - 前馈网络维度 (dim_feedforward): 内部两层全连接层中间的升维大小,默认 2048。用于提供强大的非线性拟合能力。(
int) - 丢弃率 (dropout):Dropout出现在自注意力层后、残差链接之前,也出现在前馈神经网络后、残差链接之前。默认 0.1。(
float) - 激活函数 (activation): FFN(前馈网络)内部升维后使用的非线性激活函数。支持字符串 ‘relu’(默认)或 ‘gelu’,也可以直接传入一个函数对象。(
str或Callable) - 层归一化微小值 (layer_norm_eps): 传入Layer Normalization中的一个极小值 ϵ\epsilonϵ (Epsilon),默认 1e-05。它的唯一作用是加在求方差的除法分母上,防止出现“除以 0”的数学崩溃(数值稳定性保障),绝大多数情况下永远不需要改动。(
float) - 批次优先 (batch_first): 若设为
True,则输入输出张量形状为(batch, seq, feature)。(bool) - 归一化顺序 (norm_first): 决定LayerNorm是放在Attention/FFN 之前 (Pre-Norm) 还是之后 (Post-Norm)。默认
False(Post-Norm)。在训练几十层甚至上百层的深层大模型时,工业界通常将其设为True(Pre-Norm) 以保证梯度稳定。(bool) - 偏置项开关 (bias): 决定内部的线性层 (Linear) 和层归一化 (LayerNorm) 是否学习并使用偏置项 (Bias)。默认设为 True。(
bool) - 分配设备 (device): 指定该层初始化的参数权重存放在哪个硬件设备上(如 torch.device(‘cuda:0’) 或 ‘cpu’)。(
torch.device或str或None) - 数据类型 (dtype): 指定该层参数张量的数据类型(如 torch.float32,torch.float16,torch.bfloat16)。默认 None(即使用全局默认的 float32)。(
torch.dtype或None)
理论与底层代码的对齐 (极其重要):
在底层实现中,Decoder 层经历了三步走战略(以 Post-Norm 为例):
-
带掩码的自注意力 (防作弊):
x′=LayerNorm(x+MultiHeadAttention(x,x,x,attn_mask=M))x' = \text{LayerNorm}(x + \text{MultiHeadAttention}(x, x, x, \text{attn\_mask}=M))x′=LayerNorm(x+MultiHeadAttention(x,x,x,attn_mask=M))
(注:这里的 MMM 就是的下三角矩阵
tgt_mask) -
交叉注意力 (查阅原文):
x′′=LayerNorm(x′+MultiHeadAttention(x′,m,m))x'' = \text{LayerNorm}(x' + \text{MultiHeadAttention}(x', m, m))x′′=LayerNorm(x′+MultiHeadAttention(x′,m,m))
(注:x′x'x′ 作为 Query,来自 Encoder 的 mmm (memory) 作为 Key 和 Value)
-
前馈网络:
y=LayerNorm(x′′+FFN(x′′))y = \text{LayerNorm}(x'' + \text{FFN}(x''))y=LayerNorm(x′′+FFN(x′′))
补充 (史诗级避坑警告):
双重 Mask 的地狱级混淆:
在 Decoder 中,由于既要算自注意力,又要算交叉注意力,它接受的 Mask 参数多达 4 个!最常搞混的是tgt_mask和memory_mask:
tgt_mask(自回归掩码):这是必须给的!形状是(tgt_seq_len, tgt_seq_len),用于防止偷看未来。memory_mask(记忆掩码):绝大多数情况下根本不用给(传None)!只有当你出于某种黑科技目的,想让 Decoder 在查阅原文时故意忽略原文的某些部分时才用。不要把tgt_mask错传给memory_mask!
(2). _call_ - 实例调用(执行前向传播)
作用:接收解码器当前输入和编码器的输出,执行单层 Decoder 的前向计算。
decoder_layer(tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)
参数:
- 目标输入张量 (tgt): 即
decoder_input,目前已生成的序列特征矩阵 (Tensor)。 - 记忆特征张量 (memory): 即
encoder_output,由 Encoder 层输出的原文全局特征矩阵 (Tensor)。 - 目标掩码 (tgt_mask): 作用于目标序列自注意力计算的自回归掩码(通常为下三角矩阵),必须传入,用于防止模型偷看未来词。形状必须是
(tgt_seq_len, tgt_seq_len)。注意其逻辑:通常上三角区域全为负无穷(代表掩盖/阻断未来信息),主对角线及下三角为 0(代表通行)。 (Tensor或None) - 记忆掩码 (memory_mask): 作用于交叉注意力计算的掩码,用于人为限制 Decoder 对
memory特定位置的可见性,极少使用。形状必须是(tgt_seq_len, src_seq_len)。用于人为限制交叉注意力中特定 Token 之间的互相可见性。(Tensor或None) - 目标补齐掩码 (tgt_key_padding_mask): 作用于目标序列的布尔矩阵,告诉模型忽略
tgt中为了凑长度而填充的<PAD>字符。形状必须是(batch_size, tgt_seq_len)。注意其反直觉逻辑:True代表该位置是无意义的<PAD>需要被掩盖/忽略,False代表真实有效的单词需要参与计算。 (Tensor或None) 记忆补齐掩码 (memory_key_padding_mask): 作用于原文记忆序列的布尔矩阵,告诉 Decoder 查阅时忽略memory中的<PAD>字符。形状必须是(batch_size, src_seq_len)。同理反直觉逻辑:True代表掩盖/忽略无意义字符,False代表真实有效。一般都不用,因为在encoder时memory就已经掩码过了。 (Tensor或None)
返回值:
- 成功: 返回单层解码后的特征张量,形状与输入的
tgt绝对相同 (Tensor)。
示例:
batch_size = 2
src_len = 10
tgt_len = 4
d_model = 512
# 1.准备数据与 Mask
encoder_input = torch.rand(batch_size, src_len, d_model)
decoder_input = torch.rand(batch_size, tgt_len, d_model)
tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_len)
# 2.实例化单层 Encoder 和 Decoder (开启 batch_first)
enc_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=8, batch_first=True)
dec_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=8, batch_first=True)
# 3.前向计算
# 第一步:Encoder 层提取原文特征
encoder_output = enc_layer(src=encoder_input)
print("Encoder Output Shape:", encoder_output.shape)
# 返回: torch.Size([2, 10, 512])
# 第二步:Decoder 层进行自回归特征融合
decoder_output = dec_layer(
tgt=decoder_input,
memory=encoder_output,
tgt_mask=tgt_mask
)
print("Decoder Output Shape:", decoder_output.shape)
# 返回: torch.Size([2, 4, 512])
5. nn.TransformerEncoder类 - Transformer编码器 (堆叠层)
nn.TransformerEncoder是将单一解码器堆叠后构成的解码器串
nn.TransformerEncoder包含了多个nn.TransformerEncoderLayer层的堆叠
nn.TransformerEncoder类同样是nn.Module类的子类
(1). nn.TransformerEncoder类类型
作用:
nn.TransformerEncoder 是整个 Transformer 架构的“左半边”。它并不是单层结构,而是一个“容器”,负责将我们之前定义的 nn.TransformerEncoderLayer(单层编码器)像千层饼一样堆叠 N 层。它的核心任务是将输入的源序列特征,通过多层的自注意力计算,提炼成具备极深上下文语义的最终全局特征(即 memory 或 encoder_output)。
torch.nn.TransformerEncoder(encoder_layer, num_layers, norm=None)
核心参数:
-
编码器层实例 (encoder_layer): 必须传入一个已经实例化好的
nn.TransformerEncoderLayer对象,作为堆叠的基础积木。(nn.TransformerEncoderLayer) -
层数 (num_layers): 指定要将上述积木堆叠多少层。经典的 Transformer 模型默认是 6 层。(
int) -
归一化层 (norm): (可选) 传入一个实例化好的归一化层(通常是
nn.LayerNorm),它会在经过所有N层编码器处理完之后,对最终的输出再做一次归一化。(nn.Module或None)
理论与底层代码的对齐 (极其重要):
在底层前向传播时,它本质上就是一个强大的 for 循环。源输入 XsrcX_{\text{src}}Xsrc 会被串行地穿过每一层:
H0=XsrcH_0 = X_{\text{src}}H0=Xsrc
Hi=EncoderLayeri(Hi−1),for i=1,…,NH_i = \text{EncoderLayer}_i(H_{i-1}), \quad \text{for } i = 1, \dots, NHi=EncoderLayeri(Hi−1),for i=1,…,N
Eout=Norm(HN)E_{\text{out}} = \text{Norm}(H_N)Eout=Norm(HN)
其中 EoutE_{\text{out}}Eout 就是输出的终极特征矩阵,它将作为已知条件发送给 Decoder。
补充 (史诗级避坑警告):
浅拷贝与深拷贝的参数共享陷阱:
在传入encoder_layer时,你只需要实例化“一个”单层对象传进去即可。底层会使用copy.deepcopy自动帮你复制出N个相互独立、权重不共享的层。
千万不要自作聪明地用[nn.TransformerEncoderLayer(...)] * 6这种 Python 列表语法去构建!那样会导致 6 层的权重地址完全一样(参数共享),你的深层网络瞬间退化成单层网络循环跑 6 遍!
(2). _call_ - 实例调用(执行前向传播)
作用:接收最原始的源序列张量和对应的掩码,一次性穿透所有堆叠的 Encoder 层,输出终极特征图。
我们常定义nn.TransformerEncoder类型对象叫encoder
# 实例化后,直接像调用函数一样调用该对象
encoder(src, mask=None, src_key_padding_mask=None)
参数:
- 源输入张量 (src): 也就是需要处理的原文特征矩阵。(
Tensor) 源注意力掩码 (mask): 极少使用 (通常为None),除非有特殊控制需求。形状必须是(src_seq_len, src_seq_len)。用于人为限制源序列内部特定 Token 之间的互相可见性。 (Tensor或None)- 补齐掩码 (src_key_padding_mask): 屏蔽原文中
<PAD>字符的布尔张量。形状必须是(batch_size, src_seq_len)。注意其反直觉逻辑:True代表该位置是无意义的<PAD>需要被掩盖/忽略,False代表真实有效的单词需要参与计算。 传入后,这个掩码会自动作用于内部的每一层!(Tensor或None)
返回值:
- 成功: 返回经过 NNN 层深度提取后的输出张量,形状与输入的
src绝对相同。(Tensor)
示例:
batch_size = 2
src_len = 10
d_model = 512
# 1. 准备源数据
encoder_input = torch.rand(batch_size, src_len, d_model)
# 2. 实例化单层积木 (开启 batch_first)
encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
# 3. 实例化多层容器 (堆叠 6 层)
encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
# 4. 执行多层前向计算:触发 __call__
encoder_output = encoder(src=encoder_input)
print("6层Encoder的终极输出形状:", encoder_output.shape)
# 返回: torch.Size([2, 10, 512])
6. nn.TransformerDecoder类 - Transformer解码器 (堆叠层)
nn.TransformerDecoder是将单一编码器堆叠后构成的编码器串
nn.TransformerDecoder包含了多个nn.TransformerDecoderLayer层的堆叠。
nn.TransformerDecoder类同样是nn.Module类的子类
(1). nn.TransformerDecoder类类型
作用:
nn.TransformerDecoder 是整个 Transformer 架构的“右半边”。它同样是一个容器,负责将 nn.TransformerDecoderLayer 堆叠 N 层。它的核心任务是一边进行自回归计算(看自己生成的文本),一边进行交叉注意力计算(查阅 Encoder 送来的全局特征),最终输出精准的预测特征分布。
torch.nn.TransformerDecoder(decoder_layer, num_layers, norm=None)
核心参数:
- 解码器层实例 (decoder_layer): 必须传入一个实例化的
nn.TransformerDecoderLayer。(nn.TransformerDecoderLayer) - 层数 (num_layers): 堆叠层数 (int)。通常与 Encoder 保持一致,设为 6 层。(
int) - 归一化层 (norm): (可选) 所有层执行完毕后的最终归一化操作。(
nn.Module或None)
理论与底层代码的对齐 (极其重要):
它的底层同样是一个串联的推导过程,但多了一个极度关键的常数条件——来自编码器的 EoutE_{\text{out}}Eout(即代码中的 memory):
S0=XtgtS_0 = X_{\text{tgt}}S0=Xtgt
Si=DecoderLayeri(Si−1,Eout),for i=1,…,NS_i = \text{DecoderLayer}_i(S_{i-1}, E_{\text{out}}), \quad \text{for } i = 1, \dots, NSi=DecoderLayeri(Si−1,Eout),for i=1,…,N
Dout=Norm(SN)D_{\text{out}} = \text{Norm}(S_N)Dout=Norm(SN)
每一次穿过 Decoder 层,模型都会去反复查阅那个不变的 EoutE_{\text{out}}Eout。
补充 (史诗级避坑警告):
掩码参数的全局穿透法则:
在调用堆叠的TransformerDecoder时,你只需要在最外层把 4 大 Mask(tgt_mask,memory_mask,tgt_key_padding_mask,memory_key_padding_mask)传进去一次。
底层代码会自动帮你把这些 Mask 准确无误地分发给内部的每一层(每一层的自注意力模块拿tgt_mask,交叉注意力模块拿memory相关的 Mask)。千万不要尝试自己写for循环去一层层传,直接用这个官方封装好的类是最安全、最高效的做法!
(2). _call_ - 实例调用(执行前向传播)
作用:接收已生成的目标序列和编码器的输出,穿透所有堆叠层,输出最终的预测结果。
我们常定义nn.TransformerDecoder类型对象叫decoder
# 实例化后,直接像调用函数一样调用该对象
decoder(tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None)
参数:
- 目标输入 (tgt): 即
decoder_input,目标序列特征矩阵。(Tensor) - 记忆张量 (memory): 即
encoder_output,来自前序 Encoder 容器的最终输出。(Tensor) - 目标掩码 (tgt_mask): 必须传入的下三角矩阵!形状必须是
(tgt_seq_len, tgt_seq_len)。注意其逻辑:通常上三角区域全为负无穷(代表掩盖/阻断未来信息),主对角线及下三角为 0(代表通行)。 (Tensor或None) 记忆掩码 (memory_mask): 极少使用,用于人为限制 Decoder 对 Encoder 输出特定位置的可见性。形状必须是(tgt_seq_len, src_seq_len)。(Tensor或None)- 目标补齐掩码 (tgt_key_padding_mask): 屏蔽目标序列中
<PAD>字符的布尔张量。形状必须是(batch_size, tgt_seq_len)。注意其反直觉逻辑:True代表该位置是无意义的<PAD>需要被掩盖/忽略,False代表真实有效的单词需要参与计算。 (Tensor或None) - 记忆补齐掩码 (memory_key_padding_mask): 屏蔽原文序列中
<PAD>字符的布尔张量。形状必须是(batch_size, src_seq_len)。同理反直觉逻辑:True代表掩盖/忽略无意义字符,False代表真实有效。 (Tensor或None)
返回值:
- 成功: 返回经过 NNN 层融合后的预测特征张量,形状与输入的
tgt绝对相同。(Tensor)
示例:
import torch
import torch.nn as nn
batch_size = 2
src_len = 10 # 用于模拟 memory 长度
tgt_len = 4 # 用于模拟当前生成长度
d_model = 512
# 1. 准备数据 (通常 memory 是由上一节的 encoder 算出来的)
decoder_input = torch.rand(batch_size, tgt_len, d_model) # tgt
encoder_output = torch.rand(batch_size, src_len, d_model) # memory
# 2. 准备防作弊掩码 (极其重要)
tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_len)
# 3. 实例化单层积木 (开启 batch_first)
decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
# 4. 实例化多层容器 (堆叠 6 层)
decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
# 5. 执行多层前向计算:触发 __call__
decoder_output = decoder(
tgt=decoder_input,
memory=encoder_output,
tgt_mask=tgt_mask
)
print("6层Decoder的终极输出形状:", decoder_output.shape)
# 返回: torch.Size([2, 4, 512])
7. nn.MultiheadAttention类 - 多头注意力机制
nn.MultiheadAttention类同样是nn.Module类的子类。
(1). nn.MultiheadAttention类类型
作用:
nn.MultiheadAttention: 这个模块实现了多头注意力机制,这是Transformer模型的核心组件之一。多头注意力允许模型在不同的位置同时处理来自序列不同部分的信息,这有助于捕捉序列内的复杂依赖关系。在底层,它通过将输入映射到多个头(Heads),分别计算按权重缩放后的相关度,最后再将结果拼接融合,极大地增强了模型捕捉复杂长距离依赖的能力。
torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False)
参数:
- 总特征维度 (embed_dim): 模型的输入输出特征总维度。(int)
- 多头数量 (num_heads): 并行计算注意力的头数。
embed_dim必须能被num_heads整除! (int) - 失活率 (dropout): 作用于注意力权重矩阵上的Dropout概率,默认
0.0。(float) - 批次优先 (batch_first): 决定输入张量的形状。默认是
False,这意味着 PyTorch 期望的形状是序列长度在第一维(SeqLen, Batch, Features),这常让新手抓狂;如果设为True,则符合人类直觉的(Batch, SeqLen, Features)。(bool) - 偏置项 (bias): 决定是否在 Q、K、V 的输入线性投影层以及最终的输出融合映射层中添加可学习的偏置项(Bias)。默认为
True。(bool) - 键/值单独维度 (kdim, vdim): 默认情况下 (
None),模型假定 Query、Key、Value 的输入特征维度是完全相同的,皆为embed_dim。如果是在做交叉注意力 (Cross-Attention) 时,外部传入的 Key 和 Value 的维度与 Query 不同,则需要显式指定这两个参数。(注:即便输入维度不同,底层仍会通过各自的投影矩阵将它们映射到统一的内部维度以计算点积)。(int 或 None) - 键值偏置追加 (add_bias_kv): (进阶参数) 若设为
True,会在 Key 和 Value 的序列长度维度上显式地追加一个独立的、可学习的偏置向量。默认为False。(bool) - 零注意力机制 (add_zero_attn): (进阶参数) 若设为
True,会在 Key 和 Value 的序列末尾强行追加一个全零的占位向量。这一步操作极具启发性:它在数学上为模型提供了一种合法的选项,允许注意力机制选择“什么都不关注”,在某些复杂序列任务中能提升模型的鲁棒性。默认为False。(bool)
理论与底层代码的对齐:
在数学理论中,多头注意力的计算公式为:
MultiHead(Q,K,V)=Concat(head1,...,headh)WOMultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^OMultiHead(Q,K,V)=Concat(head1,...,headh)WO
其中 headi=Attention(QWiQ,KWiK,VWiV)=softmax(qkTdk)v\text{其中 } head_i = Attention(QW_i^Q, KW_i^K, VW_i^V) = \text{softmax}\left(\frac{q k^T}{\sqrt{d_k}}\right)v其中 headi=Attention(QWiQ,KWiK,VWiV)=softmax(dkqkT)v
补充:
- WOW^OWO(Output Weight Matrix) 是一个极其关键的、可学习的参数矩阵,它的作用就相当于在多头注意力的最后接了一个
nn.Linear层。它承担着两大核心任务:- 跨头信息融合(Cross-Head Feature Fusion): 不同的注意力头(比如 head1head_1head1 负责找主谓宾,head2head_2head2 负责情感倾向)在并行计算时是各自为战的。当它们通过
Concat简单拼接在一起时,这些信息在物理上依然是隔离的。WOW^OWO 就像一个“调音台”,通过全连接的矩阵乘法,将所有头提取到的不同维度的特征进行深度的交叉与加权混合(Mix),打破头与头之间的壁垒,融合成拥有全局视野的高级特征。 - 维度对齐与还原(Dimensionality Alignment): 尽管拼接后的总维度恰好等于
embed_dim,但为了保证输出特征与输入特征处于同一个高维语义空间中(对Transformer后面顺利执行残差连接至关重要),必须通过 WOW^OWO 进行一次完整的线性变换投影。
- 跨头信息融合(Cross-Head Feature Fusion): 不同的注意力头(比如 head1head_1head1 负责找主谓宾,head2head_2head2 负责情感倾向)在并行计算时是各自为战的。当它们通过
PyTorch 底层物理动作拆解警告:
- 线性映射共享 (Q/K/V): PyTorch 底层为了追求极致的计算效率,并不会为你创建 3×h3 \times h3×h 个小的权重矩阵!相反,它会初始化一个巨大的参数矩阵
in_proj_weight(形状为3 * embed_dim, embed_dim)。- 一次性投影: 它将输入的 Query, Key, Value 一次性与这个大矩阵相乘,然后再在内存中
chunk或reshape切分成不同的头。这意味着其前向传播速度极快,但也导致如果你想手动提取某个特定头的 Q/K/V 权重矩阵,会变得非常麻烦。- 输出投影对应 (WOW^OWO 的底层化身): 在 PyTorch 底层源码中,公式里的 WOW^OWO 精确对应着
out_proj模块。它拥有自己的可学习权重out_proj.weight(形状为embed_dim, embed_dim),如果实例化时开启了偏置,还会带有out_proj.bias。
在加入dropout参数后的单头注意力完整公式为:
Attention(Q,K,V)=Dropout(softmax(QKTdk))VAttention(Q, K, V) = \text{Dropout}\left(\text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right)\right)VAttention(Q,K,V)=Dropout(softmax(dkQKT))V
(2). _call_ - 实例调用(执行注意力计算)
作用:接收 Query (查询)、Key (键)、Value (值) 张量以及各种掩码 (Mask),执行多头注意力计算,输出聚合后的上下文张量和注意力权重矩阵。
# 实例化后,直接像调用函数一样调用该对象
attn_output, attn_output_weights = multihead_attn(query, key, value,
key_padding_mask=None, attn_mask=None)
参数:
-
查询张量 (query): 检索上下文的输入向量。数据类型要求为浮点型 (
float32等)。(Tensor) -
键张量 (key): 与查询计算相关度得分的向量。数据类型要求为浮点型 (
float32等)。(Tensor) -
值张量 (value): 被加权聚合的实际特征向量。数据类型要求为浮点型 (
float32等)。(Tensor)注:
- 如果
batch_first=False(默认):以上三者形状必须是(seq_len, batch_size, input_dimensions)。 - 如果
batch_first=True:以上三者形状必须是(batch_size, seq_len, input_dimensions)。
- 如果
-
序列填充掩码 (key_padding_mask): (可选参数) 用于告诉模型哪些位置是无意义的
<PAD>占位符,不要给它们分配注意力。默认为None。(Tensor 或 None)注: 在这个掩码张量中,
True表示忽略它,False表示关注它!千万别写反了! -
注意力因果掩码 (attn_mask): (可选参数) 常用于 Decoder 中阻止信息“向后穿越” (即前面的词不能看到后面的词)。通常传入一个上三角全是负无穷、下三角全是 0 的二维矩阵。默认为
None。(Tensor 或 None)
返回值:
-
上下文输出 (attn_output):经过注意力机制加权求和后的新特征表示。形状与输入的
query完全一致。(Tensor)数学对齐: 对应着融合了各头信息后的最终结果,即公式中的:
attn_output=Concat(head1,...,headh)WOattn\_output = Concat(head_1, ..., head_h)W^Oattn_output=Concat(head1,...,headh)WO
MultiHead(Q,K,V)=Concat(head1,...,headh)WOMultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^OMultiHead(Q,K,V)=Concat(head1,...,headh)WO
其中 headi=Attention(QWiQ,KWiK,VWiV)=softmax(qkTdk)v\text{其中 } head_i = Attention(QW_i^Q, KW_i^K, VW_i^V) = \text{softmax}\left(\frac{q k^T}{\sqrt{d_k}}\right)v其中 headi=Attention(QWiQ,KWiK,VWiV)=softmax(dkqkT)v
-
注意力权重 (attn_output_weights):各个位置之间的注意力打分矩阵 (已经过 Softmax)。形状通常为
(batch_size, target_seq_len, source_seq_len)。(Tensor)数学对齐: 对应着公式中负责分配概率的权重部分(底层默认返回各头权重的平均值):
attn_output_weights=softmax(QKTdk)attn\_output\_weights = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right)attn_output_weights=softmax(dkQKT)
注意:没有乘上VVV
示例:
# 一.自注意力机制Self-Attention例子:
# 1.构造列数据(batch_first=True)
batch_size, seq_len, input_dimensions = 2, 5, 16
# 这里的x既作为 Query, 也作为Key和Value
x = torch.rand(batch_size, seq_len, input_dimensions)
# 2.实例化MultiheadAttention
# 嵌入维度16,分成4个头(每个头的特征维数为16/4 = 4)
mha = nn.MultiheadAttention(embed_dim=input_dimensions, num_heads=4, batch_first=True)
# 3.构造Padding Mask(假设第一个样本最后1个词是PAD,第二个样本最后2个词是PAD)
# 注意铁律:True 表示被遮挡,False 表示有效需要参与计算!
padding_mask = torch.tensor([
[False, False, False, False, True],
[False, False, False, True, True]
])
# 4. 执行注意力前向传播(Self-Attention则Q=K=V=x)
# 注意返回的是一个元组: (输出张量, 注意力权重张量)
attn_output, attn_weights = mha(query=x, key=x, value=x, key_padding_mask=padding_mask)
# 5. 验证结果形状
print("输出特征形状:", attn_output.shape)
# 返回: 输出特征形状: torch.Size([2, 5, 16]) -> 与输入完全相同
print("注意力权重形状:", attn_weights.shape)
# 返回: 注意力权重形状: torch.Size([2, 5, 5]) -> [Batch, Target_Seq, Source_Seq]
# 注意:不是[2, 4, 5, 5],在返回attn_weights时,默认将所有头的注意力分数取了平均值。
# 二.交叉注意力机制Cross-Attention例子:
# 1. 构造序列数据 (batch_first=True)
batch_size, input_dimensions = 2, 16
source_seq_len = 5 # Encoder 原文的序列长度 (原文共5个词)
target_seq_len = 4 # Decoder 当前的序列长度 (假设正在生成第4个词)
# Key和Value来源于Encoder的输出 (被查询的原文特征)
encoder_out = torch.rand(batch_size, source_seq_len, input_dimensions)
# Query来源于Decoder的输出 (要查询的内容)
decoder_out = torch.rand(batch_size, target_seq_len, input_dimensions)
# 2. 实例化 MultiheadAttention
# 注意:即使是交叉注意力,使用的依然是同一个类。模型通过你传入的 Q, K, V 来源不同来区分。
mha_cross = nn.MultiheadAttention(embed_dim=input_dimensions, num_heads=4, batch_first=True)
# 3. 构造Padding Mask
# 在交叉注意力中,我们要遮挡的是Encoder原文中的PAD占位符!
# 形状必须是[batch_size, source_seq_len],即 [2, 5]
encoder_padding_mask = torch.tensor([
[False, False, False, False, True], # 第一个样本原文最后1个是PAD
[False, False, False, True, True] # 第二个样本原文最后2个是PAD
])
# 4. 执行注意力前向传播 (Cross-Attention则 Q=decoder_out, K=V=encoder_out)
# 底层会一次性执行整个矩阵乘法,算出所有时间步的上下文向量,没有 for 循环!
cross_attn_output, cross_attn_weights = mha_cross(
query=decoder_out,
key=encoder_out,
value=encoder_out,
key_padding_mask=encoder_padding_mask
)
# 5. 验证结果形状
print("交叉注意力输出特征形状:", cross_attn_output.shape)
# 返回: torch.Size([2, 4, 16])
# 铁律:输出的形状永远与 Query 保持一致!(融合了 Encoder 信息的 Decoder 新特征)
print("交叉注意力权重形状:", cross_attn_weights.shape)
# 返回: torch.Size([2, 4, 5])
# 铁律:形状永远是 [Batch, Target_Seq (Decoder侧), Source_Seq (Encoder侧)]
# 含义:这 4x5 的矩阵代表了 Decoder 的 4 个词分别对 Encoder 的 5 个词分配的注意力分数。
8. nn.LayerNorm - 层标准化层
nn.LayerNorm: 层归一化(Layer Normalization)通常用在Transformer的各个子层的输出上,有助于稳定训练过程,并且提高了训练的速度和效果。
nn.LayerNorm同样是nn.Module类的子类
(1). nn.LayerNorm类类型
作用:
层标准化(Layer Normalization,简称 LN 层)是专门为了解决序列模型(如 RNN、Transformer)以及小 Batch Size 训练痛点而诞生的利器。不同于 BN 跨样本求特征的统计量,LN 的核心作用是在“每一个样本内部”,对其自身指定的特征维度求均值为 0、方差为 1 的正态分布。它彻底摆脱了对 Batch Size 的依赖,是自然语言处理 (NLP) 领域的绝对标配。
# LayerNorm 的参数签名
torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True)
核心参数:
-
归一化形状 (normalized_shape):这是唯一必填的参数 (int 或 list/torch.Size)。它告诉模型你想在哪些维度上算均值和方差。
- 如果填入一个整数(例如
256),表示只对最后一个维度(特征数为 256)进行归一化。 - 如果填入一个列表(例如
[C, H, W]),表示同时在这三个维度上拉平算均值和方差。
- 如果填入一个整数(例如
-
极小值 (eps):为了防止分母出现除以 0 的崩溃情况,加在方差里的极小常数。默认
1e-05(float)。 -
是否进行逐元素仿射变换 (elementwise_affine):类似于 BN 的
affine,如果设为True,该层将额外学习参数 γ\gammaγ(缩放)和 β\betaβ(平移)。默认为True(bool)。
补充:LN 层的史诗级避坑点(与 BN 的绝大差异)
在 LN 层中,根本没有momentum和track_running_stats这两个参数!
- 为什么没有动量? 因为 LN 永远只看“当前这一个样本”,它不需要像 BN 那样依靠一个个局部的 Batch 去估算全局数据集的总方差。
- 为什么不追踪全局统计量? 因为测试阶段(推理时),LN 依然是老老实实地计算当前输入样本的均值和方差,绝对不会去用训练时积累的什么历史均值。
核心属性 (Attributes):
实例化 LN 层后,它维护的状态比 BN 干净得多:
.weight和.bias:这就是公式中的 γ\gammaγ 和 β\betaβ!形状与你传入的normalized_shape完全一致,并且requires_grad=True(会通过反向传播被优化)。- 极其重要:LN 层没有
.running_mean和.running_var!它没有任何历史包袱。
理论与底层代码的对齐 (极其重要):
LN 底层的数学操作同样是死板的两步,但计算域完全不同:
- 标准化: 针对单个样本自身的指定维度计算均值 μL\mu_LμL 和方差 σL2\sigma^2_LσL2,将其拉回标准分布。
- 仿射变换(恢复表达能力): 利用可学习的 γ\gammaγ 和 β\betaβ 把分布进行适当拉伸和平移。
y=x−μLσL2+ϵ∗γ+βy = \frac{x - \mu_L}{\sqrt{\sigma^2_L + \epsilon}} * \gamma + \betay=σL2+ϵx−μL∗γ+β
补充:LayerNorm 的维度对齐方式(从右向左匹配)
- **匹配规则:传入的
normalized_shape必须从输入张量的最右侧(最后一个维度)**开始精确匹配。- **举例:**假设输入 NLP 数据张量形状为
(N, Seq_len, Feature)即(批次, 句子长度, 词向量维度)。- 如果设置
LayerNorm(Feature):它会对每个词的Feature维向量单独算均值为 0,方差为 1。(这是 NLP 中最常见的用法)。- 如果设置
LayerNorm([Seq_len, Feature]):它会把这一整句话所有的词、所有的特征全部拉平,算出一个唯一的均值和方差!
(2). call - 实例调用(执行前向传播)
作用:对输入数据执行层标准化。由于 LN 每次只依赖当前样本自身计算,因此在 LN 层中,train() 和 eval() 模式下的运算逻辑是完全一模一样的!(但严谨起见,整体网络代码中仍需保留模式切换逻辑)。
# 实例化后,直接作为函数调用
ln_layer(input)
参数:
- 输入张量 (input): 任意维度的张量,只要其末尾维度与
normalized_shape匹配即可。
返回值:
- 成功: 返回标准化且仿射变换后的结果,形状与输入绝对一致。
示例:
# 1.准备Transformer标准格式的假数据
# 假设有 2 个句子 (Batch=2),每个句子 5 个词 (Seq_len=5),词向量维度为 512 (d_model=512)
X = torch.randn(2, 5, 512)
# 为了演示,我们人为把第一个句子的第一个词的数值放大,模拟离群值
X[0, 0, :] = X[0, 0, :] * 100.0
# 2.实例化 LayerNorm,在 Transformer 中,我们专门对词向量维度(512)做归一化
ln = nn.LayerNorm(normalized_shape=512)
# 观察初始化的可学习参数
print("初始gamma(weight)形状:", ln.weight.shape)
# 初始gamma(weight)形状: torch.Size([512]) -> 全 1
print("初始beta(bias)形状:", ln.bias.shape)
# 初始beta(bias)形状: torch.Size([512]) -> 全 0
# 3.前向传播
z_hat = ln(X)
# 观察输出结果
print("经过 LN 处理后的输出 z_hat 形状:", z_hat.shape)
# 经过 LN 处理后的输出 z_hat 形状: torch.Size([2, 5, 512]) -> 形状完全不变
# 验证核心逻辑:LN 会在最内层的 512 维度上,对每一个词【独立】计算均值和方差!
# 验证归一化效果
print(f"样本1第1个词的均值: {z_hat[0, 0, :].mean().item():.6f}") # 均值趋近于 0
print(f"样本1第1个词的方差: {z_hat[0, 0, :].var(unbiased=False).item():.6f}") # 方差趋近于 1
print(f"样本2第5个词的均值: {z_hat[1, 4, :].mean().item():.6f}") # 均值趋近于 0
print(f"样本2第5个词的方差: {z_hat[1, 4, :].var(unbiased=False).item():.6f}") # 方差趋近于 1
# 4.测试模式
ln.eval()
test_X = torch.randn(1, 3, 512) # 来了一个新句子,长度为3
test_out = ln(test_X)
# 在eval模式下,LN依然会当场计算当前输入词向量的均值和方差,完全不使用训练时的历史数据!
# 所以eval模式和train模式对于nn.LayerNorm一样
print(f"Eval模式下新样本的均值: {test_out[0, 0, :].mean().item():.6f}") # 依旧趋近于 0
9. nn.Embedding - 词嵌入层
nn.Embedding同样是nn.Module类的子类
(1). nn.Embedding类类型
作用:
词嵌入层本质上是一个巨大的、可学习的查词字典。它的行数等于你的词汇表大小,列数等于你想要的特征维度(即 Transformer 中的 dmodeld_{model}dmodel)。当你给它一个代表单词的整数索引(Index)时,它会瞬间从底层矩阵中抽出对应的那一行向量返回给你。
torch.nn.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None, _freeze=False, device=None, dtype=None)
核心参数:
- 词汇表大小 (num_embeddings): 字典里总共有多少个词。输入张量中的最大索引值必须严格小于这个数。(
int) - 词向量维度 (embedding_dim): 每个词要被转换成多长的向量(例如 Transformer 经典的 512 维)。(
int) - 填充索引 (padding_idx): (极其重要) 指定词汇表中哪个索引代表
<PAD>(填充符)。设定后,该索引对应的权重向量会被硬编码为全 0,并且在训练过程中绝不会更新(即梯度为 0)。(int或None)
进阶参数 (日常较少使用):
- 最大范数 (max_norm): 如果某个词的向量范数(长度)超过了这个值,会将其重新缩放(Renorm)回这个最大值。(
float或None) - 稀疏梯度 (sparse): 若设为
True,则权重矩阵的梯度将变为稀疏张量。在词汇表极其巨大(如几百万)时可以优化内存,但绝大多数优化器(如 Adam)不支持稀疏梯度,所以通常保持False。(bool)
核心属性 (Attributes):
.weight: 这是 Embedding 层唯一的也是最核心的参数!它就是一个形状为(num_embeddings, embedding_dim)的二维矩阵,默认初始化为标准正态分布,并且requires_grad=True(除padding_idx对应的行外)。
理论与底层代码的对齐 (极其重要):
在早期的数学推导中,将离散的词转为稠密向量,是通过将词先转为独热编码(One-Hot Encoding),然后再乘以一个权重矩阵来实现的:
Vout=Xone-hot×WembedV_{\text{out}} = X_{\text{one-hot}} \times W_{\text{embed}}Vout=Xone-hot×Wembed
但在 PyTorch 底层,如果真的先生成 One-Hot 矩阵再去相乘,会极其浪费内存且计算缓慢。因此,nn.Embedding 彻底抛弃了矩阵乘法,底层直接由 C++ 实现为 O(1)O(1)O(1) 复杂度的数组索引寻址(Index Lookup):
Vout=Wembed[Index]V_{\text{out}} = W_{\text{embed}}[\text{Index}]Vout=Wembed[Index]
补充 (史诗级避坑警告):
Embedding 的三大经典报错陷阱:
- 越界报错 (IndexError: index out of range):
如果你的num_embeddings设为 10000(合法索引是 0 到 9999),但你输入的张量里不小心混进了一个 10000,程序会在前向传播时当场崩溃!- 数据类型报错 (RuntimeError: expected scalar type Long):
nn.Embedding是用来查字典的,字典的页码只能是整数!输入给 Embedding 层的张量数据类型必须是torch.LongTensor(即torch.int64),如果你传了float32,必然报错。- 忘记处理 PAD 导致噪音干扰:
如果不设置padding_idx,模型会把<PAD>当作一个正常的词去学习它的语义,这会严重污染上下文特征。务必在初始化时将其指定!
(2). call - 实例调用(执行前向传播)
作用:接收由整数词汇索引组成的序列张量,将其映射为高维浮点数特征张量。
# 实例化后,直接作为函数调用
embedding_layer(input)
参数:
- 输入张量 (input): 包含离散词索引的张量。形状任意(通常是
(batch_size, seq_len)),数据类型必须是torch.long。
返回值:
- 成功: 返回映射后的高维特征张量。它的形状是在输入张量的最后追加一个
embedding_dim维度。即输入(B, S),输出必然是(B, S, E)。数据类型为浮点型 (torch.float32等)。
尝试一段真实的文字作为输入
从一段话变为能够输入到embedding层的信息,还需经过以下的步骤:
- 分词
- 制作词汇表(为每个单独的词标上唯一的索引)
- 依据词汇表编码(把文字转成数字)
- 填充/裁剪编码后序列(短句子加上
<PAD>,长句子裁剪) - 所有的编码好的句子放进nn.Embedding层
示例:
# 1. 假设我们的微型词汇表只有 5 个词:
# 0: <PAD>, 1: "我", 2: "爱", 3: "PyTorch", 4: "Transformer"
vocab_size = 5
d_model = 16 # 为了方便展示,我们把维度设为16而不是512
pad_index = 0
# 2. 实例化Embedding层,并郑重声明0号索引是PAD
embed_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model, padding_idx=pad_index)
# 观察底层核心矩阵 weight
print("底层的字典矩阵形状:", embed_layer.weight.shape)
# 底层的字典矩阵形状: torch.Size([5, 16])
# 验证 padding_idx 的威力:第0行被强行锁定为全0
print("PAD对应的权重(第0行):\n", embed_layer.weight[0])
# PAD对应的权重(第0行):
# tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], grad_fn=<SelectBackward0>)
# 3. 构造一批假数据 (Batch Size = 2, Seq Len = 4)
# 注意:数据类型必须是 int64 (torch.long)
# 样本1: "我 爱 PyTorch <PAD>" -> [1, 2, 3, 0]
# 样本2: "我 爱 Transformer <PAD>" -> [1, 2, 4, 0]
input_indices = torch.tensor([
[1, 2, 3, 0],
[1, 2, 4, 0]
], dtype=torch.long)
print("\n输入张量形状:", input_indices.shape)
# 输入张量形状: torch.Size([2, 4])
# 4. 前向传播:执行查表操作
embedded_features = embed_layer(input_indices)
# 观察输出结果
print("经过 Embedding 处理后的输出形状:", embedded_features.shape)
# 经过 Embedding 处理后的输出形状: torch.Size([2, 4, 16]) -> 完美符合 Transformer 的 (Batch, Seq_len, d_model)
# 验证映射的准确性:
# 样本1 的第一个词(1:"我") 和 样本2 的第一个词(1:"我"),映射出的向量应该是完全一模一样的!
print("\n样本1第一个词的特征向量和样本2第一个词的特征向量是否完全相同?")
print(torch.equal(embedded_features[0, 0, :], embedded_features[1, 0, :])) # 返回: True
# 验证 PAD 的输出:
# 输入序列的最后一个词都是 0 (<PAD>),它查出来的特征必须全是 0
print("\n样本1最后一个词(PAD)输出的特征:\n", embedded_features[0, 3, :])
# 返回: tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], grad_fn=<SelectBackward0>)
10. nn.Transformer.generate_square_subsequent_mask() - 生成自回归掩码
作用:生成一个正方形的下三角注意力掩码(Attention Mask),用于在序列生成任务中(如 GPT、机器翻译的解码阶段)屏蔽未来信息,确保模型在预测第 t 个位置时,只能看到 t 及之前的位置,防止“信息穿越”。这个函数很少用
具体来说,该函数创建了一个方阵,其中对角线及其以下的元素为0(表示可以“看到”这些位置的元素),其余元素为负无穷大(在softmax之前应用,表示位置被屏蔽,不应该有注意力权重)。
torch.nn.Transformer.generate_square_subsequent_mask(sz, device=None, dtype=None)
参数:
- 序列长度 (sz): 生成的掩码矩阵的大小,通常严格等于输入目标序列(Target Sequence)的长度 (int)。
- 设备 (device): 掩码张量所在的设备,便于与输入张量保持在同一个 GPU/CPU 上 (torch.device 或 None)。
- 数据类型 (dtype): 掩码张量的数据类型 (torch.dtype 或 None)。建议显式指定为
torch.float32等浮点型。如果在 PyTorch 较新版本中不指定,默认会返回带有-inf和0.0的浮点张量。
返回值:
- 成功: 返回一个形状为
(sz, sz)的二维张量 (Tensor)。
本质就是构造一个掩码矩阵 MMM:
Mi,j={0,if j≤i−∞,if j>iM_{i, j} = \begin{cases} 0, & \text{if } j \le i \\ -\infty, & \text{if } j > i \end{cases}Mi,j={0,−∞,if j≤iif j>i
其中:
- iii 表示当前 Query(查询,代表当前正在解码的时刻)所在的索引。
- jjj 表示 Key/Value(被查询目标,代表序列中提供信息的时刻)所在的索引。
示例:
# 生成序列长度为 4 的自回归掩码矩阵
sz = 4
mask = nn.Transformer.generate_square_subsequent_mask(sz)
print(mask)
# tensor([[0., -inf, -inf, -inf],
# [0., 0., -inf, -inf],
# [0., 0., 0., -inf],
# [0., 0., 0., 0.]])
# 典型使用场景:放入TransformerDecoder中
# 假设batch size为2,目标输入序列tgt长度为4,特征维度为 512
tgt = torch.rand(2, 4, 512)
# Encoder的输出memory,假设batch size为2,序列长度为10,特征维度为512
memory = torch.rand(2, 10, 512)
decoder_layer = nn.TransformerDecoderLayer(
d_model=512,
nhead=8,
batch_first=True
)
transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
# 将mask作为tgt_mask传入,确保解码第t个词时,只能看到0到t的信息
output = transformer_decoder(tgt, memory, tgt_mask=mask)
print(f"\nOutput shape: {output.shape}")
# 预期输出: Output shape: torch.Size([2, 4, 512])
1336

被折叠的 条评论
为什么被折叠?



