基于Transformer的机器翻译——模型篇

1.模型结构

本案例整体采用transformer论文中提出的结构，部分设置做了调整。transformer网络结构介绍可参考博客——入门级别的Transformer模型介绍，这里着重介绍其代码实现。
模型的整体结构，包括词嵌入层，位置编码，编码器，解码器、输出层部分。

2.词嵌入层

词嵌入层用于将token转化为词向量，该层可直接调用nn模块中的Embedding方法。该方法主要包括两个参数，分别表示词表的大小（vocab_size）和词嵌入的维度(emb_size)，同时为了训练更稳定，加入了缩放因子$\sqrt{d_k}$，代码如下：

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        # 词嵌入层：将词索引映射到emb_size维的向量
        self.embedding = nn.Embedding(vocab_size, emb_size)
        # 记录嵌入维度（用于缩放）
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        # 将词索引转换为词向量，并乘以√emb_size（缩放，稳定梯度）
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

3.位置编码

位置编码层用于给序列添加位置信息，解决自注意力机制无法感知序列顺序的问题。公式为：
$$PE(pos,2i)=sin(\frac{pos}{1000\frac{2i}{d}})$$
$$PE(pos,2i+1)=cos(\frac{pos}{1000\frac{2i}{d}})$$
代码表示如下：

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        # 计算位置编码的衰减因子（控制正弦/余弦函数的频率）
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        # 位置索引（0到maxlen-1）
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        # 初始化位置编码矩阵（形状：[maxlen, emb_size]）
        pos_embedding = torch.zeros((maxlen, emb_size))
        # 偶数列用正弦函数填充（pos * den）
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        # 奇数列用余弦函数填充（pos * den）
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 调整维度（添加批次维度，便于与词嵌入向量相加）
        pos_embedding = pos_embedding.unsqueeze(-2)

        # Dropout层（正则化，防止过拟合）
        self.dropout = nn.Dropout(dropout)
        # 注册为缓冲区（模型保存/加载时自动处理）
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 将词嵌入向量与位置编码相加，并应用Dropout
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])

4.编码器

由于编码器部分是通过堆叠多个子编码器层所构成的，子编码器包括：多头自注意力层、残差连接与归一化、前馈网络三部分，该部分代码全部被封装成TransformerEncoderLayer函数中，使用时只需要传递相应超参数即可，如词嵌入维度、多头注意力的头数、前馈网络的隐含层维度，代码实现为：

# 定义编码器层（单头注意力→多头注意力→前馈网络）
encoder_layer = TransformerEncoderLayer(
    d_model=emb_size,       # 输入特征维度（与词嵌入维度一致）
    nhead=NHEAD,            # 多头注意力的头数
    dim_feedforward=dim_feedforward  # 前馈网络隐藏层维度
)
# 堆叠多层编码器层形成完整编码器
self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)

5.解码器

解码器同编码器类似，代码可以表述为：

# 定义解码器层（掩码多头注意力→编码器-解码器多头注意力→前馈网络）
decoder_layer = TransformerDecoderLayer(
    d_model=emb_size,       # 输入特征维度（与词嵌入维度一致）
    nhead=NHEAD,            # 多头注意力头数（与编码器一致）
    dim_feedforward=dim_feedforward  # 前馈网络隐藏层维度
)
# 堆叠多层解码器层形成完整解码器
self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

6.输出层

输出通过线性层得到每个单词的得分，可直接通过Linear层直接实现。

7.大体代码

基于上述介绍，完整代码如下：

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)


class Seq2SeqTransformer(nn.Module):
    """
    基于Transformer的序列到序列翻译模型（日中机器翻译核心模块）
    包含编码器（处理源语言序列）和解码器（生成目标语言序列）
    """
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward: int = 512, dropout: float = 0.1):
        """
        初始化Transformer模型参数和组件
        :param num_encoder_layers: 编码器层数（论文中通常为6，此处根据计算资源调整）
        :param num_decoder_layers: 解码器层数（与编码器层数一致）
        :param emb_size: 词嵌入维度（对应Transformer的d_model，需与多头注意力维度匹配）
        :param src_vocab_size: 源语言（日语）词表大小
        :param tgt_vocab_size: 目标语言（中文）词表大小
        :param dim_feedforward: 前馈网络隐藏层维度（通常为4*d_model）
        :param dropout:  dropout概率（用于正则化，防止过拟合）
        """
        super(Seq2SeqTransformer, self).__init__()
        # 定义编码器层（单头注意力→多头注意力→前馈网络）
        encoder_layer = TransformerEncoderLayer(
            d_model=emb_size,       # 输入特征维度（与词嵌入维度一致）
            nhead=NHEAD,            # 多头注意力的头数（需满足 emb_size % nhead == 0）
            dim_feedforward=dim_feedforward  # 前馈网络隐藏层维度
        )
        # 堆叠多层编码器层形成完整编码器
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)

        # 定义解码器层（掩码多头注意力→编码器-解码器多头注意力→前馈网络）
        decoder_layer = TransformerDecoderLayer(
            d_model=emb_size,       # 输入特征维度（与词嵌入维度一致）
            nhead=NHEAD,            # 多头注意力头数（与编码器一致）
            dim_feedforward=dim_feedforward  # 前馈网络隐藏层维度
        )
        # 堆叠多层解码器层形成完整解码器
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)

        # 生成器：将解码器输出映射到目标词表（预测每个位置的目标词）
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        # 源语言词嵌入层（将词索引转换为连续向量）
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        # 目标语言词嵌入层（与源语言共享嵌入层可提升效果，此处未共享）
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        # 位置编码层（注入序列位置信息，解决Transformer的位置无关性）
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)

    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        """
        前向传播（训练时使用教师强制，输入完整目标序列）
        :param src: 源语言序列张量（形状：[seq_len, batch_size]）
        :param trg: 目标语言序列张量（形状：[seq_len, batch_size]）
        :param src_mask: 源序列注意力掩码（形状：[seq_len, seq_len]，全0表示无掩码）
        :param tgt_mask: 目标序列掩码（下三角掩码，防止关注未来词）
        :param src_padding_mask: 源序列填充掩码（标记<pad>位置，形状：[batch_size, seq_len]）
        :param tgt_padding_mask: 目标序列填充掩码（标记<pad>位置，形状：[batch_size, seq_len]）
        :param memory_key_padding_mask: 编码器输出的填充掩码（与src_padding_mask一致）
        :return: 目标序列的词表概率分布（形状：[seq_len, batch_size, tgt_vocab_size]）
        """
        # 源序列处理：词嵌入 + 位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        # 目标序列处理：词嵌入 + 位置编码（训练时使用教师强制，输入完整目标序列）
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        # 编码器处理源序列，生成记忆向量（memory）
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        # 解码器利用记忆向量生成目标序列
        outs = self.transformer_decoder(
            tgt_emb,                # 目标序列嵌入（含位置信息）
            memory,                 # 编码器输出的记忆向量
            tgt_mask,               # 目标序列掩码（防止未来词）
            None,                   # 编码器-解码器注意力掩码（此处未使用）
            tgt_padding_mask,       # 目标序列填充掩码（忽略<pad>）
            memory_key_padding_mask # 记忆向量填充掩码（与源序列填充掩码一致）
        )
        # 通过生成器输出目标词表的概率分布
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        """
        编码源序列（推理时单独调用，生成编码器记忆向量）
        :param src: 源语言序列张量（形状：[seq_len, batch_size]）
        :param src_mask: 源序列注意力掩码（形状：[seq_len, seq_len]）
        :return: 编码器输出的记忆向量（形状：[seq_len, batch_size, emb_size]）
        """
        return self.transformer_encoder(
            self.positional_encoding(self.src_tok_emb(src)),  # 源序列嵌入+位置编码
            src_mask  # 源序列注意力掩码
        )

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        """
        解码目标序列（推理时逐步生成目标词）
        :param tgt: 当前已生成的目标序列前缀（形状：[current_seq_len, batch_size]）
        :param memory: 编码器输出的记忆向量（形状：[seq_len, batch_size, emb_size]）
        :param tgt_mask: 目标序列掩码（下三角掩码，防止关注未来词）
        :return: 解码器输出（形状：[current_seq_len, batch_size, emb_size]）
        """
        return self.transformer_decoder(
            self.positional_encoding(self.tgt_tok_emb(tgt)),  # 目标前缀嵌入+位置编码
            memory,  # 编码器记忆向量
            tgt_mask  # 目标前缀掩码（仅允许关注已生成部分）
        )
class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        # 计算位置编码的衰减因子（控制正弦/余弦函数的频率）
        den = torch.exp(- torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        # 位置索引（0到maxlen-1）
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        # 初始化位置编码矩阵（形状：[maxlen, emb_size]）
        pos_embedding = torch.zeros((maxlen, emb_size))
        # 偶数列用正弦函数填充（pos * den）
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        # 奇数列用余弦函数填充（pos * den）
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 调整维度（添加批次维度，便于与词嵌入向量相加）
        pos_embedding = pos_embedding.unsqueeze(-2)

        # Dropout层（正则化，防止过拟合）
        self.dropout = nn.Dropout(dropout)
        # 注册为缓冲区（模型保存/加载时自动处理）
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 将词嵌入向量与位置编码相加，并应用Dropout
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        # 词嵌入层：将词索引映射到emb_size维的向量
        self.embedding = nn.Embedding(vocab_size, emb_size)
        # 记录嵌入维度（用于缩放）
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        # 将词索引转换为词向量，并乘以√emb_size（缩放，稳定梯度）
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

结语

至此，模型已完成搭建，后续博客将继续介绍模型训练部分的内容，希望本篇博客能够对你理解transformer有所帮助！