7.4 API使用

7.4.1 概述

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

7.4.2 核心类

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

• nn.Transformer

封装了完整的 Transformer架构，由编码器和解码器组成。作为顶层接口，适用于需要同时使用编码器和解码器的任务，如机器翻译。支持用户通过参数自定义层数、注意力头数、隐藏维度等模型结构。

• nn.TransformerEncoder

实现了Transformer编码器结构，由多个编码器层的堆叠而成，用于将输入序列编码为上下文相关的表示。

• nn.TransformerDecoder

实现了Transformer解码器结构，由多个解码器层堆叠而成，用于基于编码结果逐步生成目标序列。

• nn.TransformerEncoderLayer

实现了单个编码器层结构，包含一个多头自注意力子层和一个前馈神经网络子层，两者均带有残差连接和 LayerNorm。

• nn.TransformerDecoderLayer

实现了单个解码器层结构，包含自注意力、编码器-解码器注意力、前馈子层，同样配有残差连接和 LayerNorm。

7.4.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)

各参数含义如下：

• 基础核心参数

参数名	默认值	说明
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。

高级参数

参数名	默认值	说明
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
)

7.4.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
)

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

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

• 输入

输出