大模型学习15 - npl自然语言处理16 - Transformer模型API(PyTorch)

1 概述

PyTorch 提供了对 Transformer 的官方实现，该模块封装了完整的编码器-解码器结构，可直接应用于机器翻译、文本生成等典型的序列建模任务。

2 核心类

PyTorch 中的 Transformer 模块由以下几个核心类构成：

• nn.Transformer
  封装了完整的 Transformer架构，由编码器和解码器组成。作为顶层接口，适用于需要同时使用编码器和解码器的任务，如机器翻译。支持用户通过参数自定义层数、注意力头数、隐藏维度等模型结构。
• nn.TransformerEncoder
  实现了Transformer编码器结构，由多个编码器层的堆叠而成，用于将输入序列编码为上下文相关的表示。
• nn.TransformerDecoder
  实现了Transformer解码器结构，由多个解码器层堆叠而成，用于基于编码结果逐步生成目标序列。
• nn.TransformerEncoderLayer
  实现了单个编码器层结构，包含一个多头自注意力子层和一个前馈神经网络子层，两者均带有残差连接和 LayerNorm。
• nn.TransformerDecoderLayer
  实现了单个解码器层结构，包含自注意力、编码器-解码器注意力、前馈子层，同样配有残差连接和 LayerNorm。

3 Transformer 构造参数

构造Transformer模型所需参数如下：

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)

各参数含义如下：

3.1 基础核心参数

参数名	默认值	说明
d_model	512	编码器和解码器输入的特征维度，即每个 token 的表示向量长度。
nhead	8	多头注意力中的头数，要求 d_model 能被整除。
num_encoder_layers	6	编码器中堆叠的层数，决定模型深度。
num_decoder_layers	6	解码器中堆叠的层数，结构与编码器相似。
dim_feedforward	2048	前馈神经网络中的隐藏层维度，默认是 d_model 的 4 倍。
dropout	0.1	Dropout 概率，用于缓解过拟合，作用于注意力层和前馈层。
batch_first	False	若设为 True，输入和输出张量的维度格式为 (batch_size, seq_len, d_model)；否则为 (seq_len, batch_size, d_model)。推荐设为 True。
device	None	模型参数所使用的设备，如 'cuda' 或 'cpu'。
dtype	None	模型参数的数据类型，如 torch.float32 或 torch.float16。

3.2 高级参数

参数名	默认值	说明
activation	'relu'	编码器/解码器中间层的激活函数，支持字符串（'relu' 或 'gelu'）或自定义函数。
custom_encoder	None	自定义编码器模块。
custom_decoder	None	自定义解码器模块。
layer_norm_eps	1e-5	LayerNorm 层中的 epsilon 值，用于防止除零。
norm_first	False	若为 True，编码器和解码器层将在注意力或前馈操作之前执行 LayerNorm。
bias	True	若为 False，则所有 Linear 和 LayerNorm 层不使用偏置项。

示例代码如下：

from torch import nn

# 初始化 Transformer
transformer = nn.Transformer(
    d_model=512, 
    nhead=8, 
    num_encoder_layers=6, 
    num_decoder_layers=6, 
    batch_first=True
)

4 Transformer.forward

nn.Transformer 封装了完整的前向传播逻辑，其 forward() 方法定义了编码器解码器的执行流程。

该函数接收 源语言序列 （src_sequence，编码器输入）和目标语言序列 （tgt_sequence，解码器输入）作为输入，以解码器预测结果作为输出。

如下图所示：

示例代码如下：

output = transformer(
    src=src_emb,
    tgt=tgt_emb,
    src_key_padding_mask=src_pad_mask,
    tgt_key_padding_mask=tgt_pad_mask,
    tgt_mask=tgt_mask,
    memory_key_padding_mask=src_pad_mask
)

假设当前的源序列 和 目标序列为：

则具体的输入和输出内容如下

4.1 输入

1. batch_size（批次维度，斜向）
   一批里面有多少条样本。

例：batch_size=3，一次性并行处理 3 个句子。

2. src_len（源序列长度，向下）
   单个句子的 token 数量。
   图中黄色样本句子：我 爱 你，后面跟着 <pad> 填充符。

不同原始句子长短不一，batch 内统一 padding 补齐到相同 src_len，<pad>是空占位符，后面会用padding mask屏蔽掉，不让注意力看到 pad 位置。

3. d_model（模型维度，向右）
   每一个 token 对应的向量维度，也就是词嵌入向量的长度。Transformer 论文默认d_model=512。

每一个字（我/爱/<pad>）都被表示成一个长度为d_model的向量。

• src

  源序列的嵌入表示，通常由词向量与位置编码相加得到，作为编码器的输入。其形状为 (batch_size, src_len, d_model)

  

• tgt

  目标序列的嵌入表示，通常由词向量与位置编码相加得到，作为解码器的输入。其形状为 (batch_size, tgt_len, d_model)。

  

• src_key_padding_mask

  用于编码器中的自注意力机制，用以屏蔽源序列中填充（）的位置，防止模型在计算注意力时关注无意义的内容。

  其张量形状为 (batch_size, src_len)，其中值为 True 的位置表示应被忽略（即不参与注意力计算）。

  

• tgt_key_padding_mask

  用于解码器中的自注意力机制，用以屏蔽目标序列中填充（）的位置，防止模型在计算注意力时关注无意义的内容。其张量形状为 (batch_size, tgt_len)，其中值为 True 的位置表示应被忽略（即不参与注意力计算）。

  

• tgt_mask

  用于解码器的自注意力机制，常用于训练阶段的自回归任务，防止模型关注当前位置之后的 token，避免信息泄露。是一个形状为 (tgt_len, tgt_len) 的上三角矩阵，类型为 float，遮挡位置为 -inf，其余为 0。也支持 bool 类型（True 表示遮挡），内部会自动转换为加性掩码。

  

• memory_key_padding_mask

  用于解码器的交叉注意力机制，屏蔽编码器输出中的 位置，防止解码器关注源序列中的无效 token。形状为 (batch_size, src_len)，值为 True 的位置将被忽略。通常与 src_key_padding_mask 相同。

  

4.2 输出

output

解码器输出的隐藏状态序列，形状为 (batch_size, tgt_len, d_model)。表示目标序列中每个位置的上下文表示，通常会送入线性层和 softmax，用于生成词表上的预测概率。

5 Transformer.encoder

nn.Transformer 模块中包含一个编码器部分，可通过属性 transformer.encoder 访问，其本质是一个 nn.TransformerEncoder 实例。通过其 forward 方法，可以对源序列进行编码，提取上下文相关的语义表示。

如下图所示：

示例代码如下：

from torch import nn

# 初始化 Transformer
transformer = nn.Transformer(
    d_model=512, nhead=8,
    num_encoder_layers=6, num_decoder_layers=6,
    batch_first=True
)

# 调用编码器
memory = transformer.encoder(
    src=src_emb, 
    src_key_padding_mask=src_pad_mask
)

假如源序列为：

5.1 输入

batch_size（批次维度，斜向）

一批里面有多少条样本。

src_len（源序列长度，向下）

单个句子的 token 数量。

图中黄色样本句子：我 爱 你，后面跟着 <pad> 填充符。

d_model（模型维度，向右）

每一个 token 对应的向量维度，也就是词嵌入向量的长度。Transformer 论文默认d_model=512。

• src
  源序列的嵌入表示，通常由词向量与位置编码相加得到，作为编码器的输入。其形状为 (batch_size, src_len, d_model)
  
• src_key_padding_mask
  用于编码器中的自注意力机制，用以屏蔽源序列中填充（）的位置，防止模型在计算注意力时关注无意义的内容。其张量形状为 (batch_size, src_len)，其中值为 True 的位置表示应被忽略（即不参与注意力计算）。
  

5.2 输出

memory

编码器的输出表示，包含每个 token 的上下文语义信息，作为解码器的输入。其形状为 (batch_size, src_len, d_model)。

6 Transformer.decoder

nn.Transformer 模块中包含一个解码器部分，可通过属性 transformer.decoder 访问，其本质是一个 nn.TransformerDecoder 实例。通过其 forward 方法，可以基于编码器的输出（memory）和目标序列的嵌入表示，逐步生成目标序列中的各个 token。

如下图所示：

示例代码如下：

from torch import nn

# 初始化 Transformer
transformer = nn.Transformer(
    d_model=512, nhead=8,
    num_encoder_layers=6, num_decoder_layers=6,
    batch_first=True
)

# 调用编码器
memory = transformer.encoder(
    src=src_emb, 
    src_key_padding_mask=src_pad_mask
)

# 调用解码器（逐步生成）
output = transformer.decoder(
    tgt=tgt_emb,
    memory=memory,
    tgt_mask=tgt_mask,
    tgt_key_padding_mask=tgt_pad_mask,
    memory_key_padding_mask=src_pad_mask
)

若源序列和目标序列为：

则具体输出输出为：

6.1 输入

• tgt
  目标序列的嵌入表示，通常由词向量与位置编码相加得到，作为解码器的输入。其形状为 (batch_size, tgt_len, d_model)。
  
• memory
  编码器的输出表示，包含源序列每个 token 的上下文语义信息，作为解码器的输入。其形状为 (batch_size, src_len, d_model)。
  
• tgt_mask
  用于解码器的自注意力机制，常用于训练阶段的自回归任务，防止模型关注当前位置之后的 token，避免信息泄露。是一个形状为 (tgt_len, tgt_len) 的上三角矩阵，类型为 float，遮挡位置为 -inf，其余为 0。也支持 bool 类型（True 表示遮挡），内部会自动转换为加性掩码。
  
• tgt_key_padding_mask
  用于解码器中的自注意力机制，用以屏蔽源序列中填充（）的位置，防止模型在计算注意力时关注无意义的内容。其张量形状为 (batch_size, tgt_len)，其中值为 True 的位置表示应被忽略（即不参与注意力计算）。
  
• memory_key_padding_mask
  用于解码器的交叉注意力机制，屏蔽编码器输出中的 位置，防止解码器关注源序列中的无效 token。形状为 (batch_size, src_len)，值为 True 的位置将被忽略。通常与 src_key_padding_mask 相同。
  

6.2 输出

output

解码器输出的隐藏状态序列，形状为 (batch_size, tgt_len, d_model)。表示目标序列中每个位置的上下文表示，通常会送入线性层和 softmax，用于生成词表上的预测概率。