【值得收藏】手把手教你用PyTorch构建Transformer英汉翻译系统，从训练到推理

本文演示一个简化版的英汉机器翻译系统，完整呈现其训练与推理流程。我们将使用 PyTorch 逐步搭建整个流程，重点突出核心逻辑，代码简洁易懂，便于理解。

需要说明的是，本文的目标并非构建高精度的翻译系统，而是通过最小样本和最简配置，演示 Transformer 从数据准备到模型构建，再到生成翻译结果的完整闭环。

一、准备数据：构建中英文平行语料

在训练 Transformer 模型之前，我们需要准备一组中英文对照的数据，并实现基本的编码器（Tokenizer），将文本转换为模型可读的数字序列。

（一）构建中英文句子对

我们先手动构建一个小型中英文平行语料，数据仅用于结构验证，并非用于泛化或真实训练。

复制代码

raw_data = [    ("I love machine learning", "我 喜欢 机器 学习"),    ("Deep learning is powerful", "深度 学习 很 强大"),    ("Transformer changed everything", "Transformer 改变 了 一切"),]

这里每个元组表示一组输入输出对，左边是英文句子，右边是中文翻译（已用空格分词，便于处理）。

（二）构建词表

接下来我们为每种语言分别构建词表（Vocabulary），包括一些特殊标记：

复制代码

def build_vocab(sentences, min_freq=1):    from collections import Counter    counter = Counter()    for s in sentences:        counter.update(s.split())    vocab = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "<unk>": 3}    for word, freq in counter.items():        if freq >= min_freq and word not in vocab:            vocab[word] = len(vocab)    return vocab
# 拆分源语言和目标语言句子src_sentences = [s[0].lower() for s in raw_data]  # 英文转小写tgt_sentences = [s[1] for s in raw_data]          # 中文已分词
# 分别构建词表src_vocab = build_vocab(src_sentences)tgt_vocab = build_vocab(tgt_sentences)

每个词表中包含了 <pad>（填充）、<bos>（句首）、<eos>（句尾）、<unk>（未知词）等特殊 token，以便后续处理。

（三）编码器（Tokenizer）

我们实现一个简单的编码器，将句子转为固定长度的 token ID 序列，并自动添加句首、句尾标记。

复制代码

def encode(sentence, vocab, max_len=10):    tokens = sentence.lower().split()    ids = [vocab.get(tok, vocab["<unk>"]) for tok in tokens]    ids = [vocab["<bos>"]] + ids + [vocab["<eos>"]]    # 补齐或截断    if len(ids) < max_len:        ids += [vocab["<pad>"]] * (max_len - len(ids))    else:        ids = ids[:max_len]    return ids

（四）编码中英文数据

将原始句子编码为固定长度的整数序列，作为模型的输入和目标输出：

复制代码

input_seqs = [encode(s[0], src_vocab, max_len=10) for s in raw_data]target_seqs = [encode(s[1], tgt_vocab, max_len=10) for s in raw_data]

至此，我们已经得到了：

input_seqs：模型的输入序列（英文，编码后）
target_seqs：模型的目标输出（中文，编码后）

它们都是长度为 10 的 token ID 列表，可以直接用于 Transformer 模型的训练。

二、定义 Transformer 模型结构

我们将使用 PyTorch 搭建完整的 Transformer 模型，包括嵌入层、位置编码、注意力机制、前馈网络、编码器和解码器。

（1）位置编码

在 Transformer 中，由于没有像 RNN 或 CNN 那样的顺序结构，模型本身无法感知序列中各个单词的先后顺序。因此，我们需要一种机制来显式地注入位置信息，使模型能够理解词与词之间的相对或绝对位置------这正是位置编码的作用。

下面是我们用 PyTorch 实现的位置编码模块：

复制代码

class PositionalEncoding(nn.Module):    def __init__(self, d_model, max_len=5000):        super().__init__()                # 创建一个形状为 (max_len, d_model) 的全 0 矩阵        pe = torch.zeros(max_len, d_model)        # 生成每个位置的索引：0, 1, 2, ..., max_len-1        pos = torch.arange(0, max_len).unsqueeze(1)        # 计算每个维度对应的位置编码频率        div = torch.exp(torch.arange(0, d_model, 2) * -(torch.log(torch.tensor(10000.0)) / d_model))        # 将正弦函数应用于偶数维度        pe[:, 0::2] = torch.sin(pos * div)        # 将余弦函数应用于奇数维度        pe[:, 1::2] = torch.cos(pos * div)        # 增加一个 batch 维度，方便与输入相加        self.pe = pe.unsqueeze(0)        def forward(self, x):        # 将前 seq_len 个位置编码加到输入上        return x + self.pe[:, :x.size(1)].to(x.device)

位置编码不会直接告诉模型"第几个词是主语"或"谁在前谁在后"，而是通过数学函数为每个位置生成一组独特的数值向量，这些向量中包含了能反映位置差异的数学规律，模型通过学习这些规律，逐渐建立起对词序的理解。

（2）多头注意力机制

多头注意力机制是 Transformer 的核心模块之一。它的主要作用是：

并行从多个子空间（head）学习不同的注意力表示；
提高模型对不同语义关系的理解能力；
在序列的不同位置捕捉更丰富、更细腻的上下文依赖。

在这个模块中，输入的每个词向量会被映射为**查询（Q）、键（K）、值（V）**三组向量，通过它们计算注意力权重，再结合值向量得到最终的注意力输出。

复制代码

class MultiHeadAttention(nn.Module):    def __init__(self, d_model, num_heads):        super().__init__()        self.num_heads = num_heads        self.head_dim = d_model // num_heads        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"                # 为 Q（查询）、K（键）、V（值）分别定义线性变换层        self.q_proj = nn.Linear(d_model, d_model)        self.k_proj = nn.Linear(d_model, d_model)        self.v_proj = nn.Linear(d_model, d_model)                # 最后的线性层，用于整合所有头的输出        self.out_proj = nn.Linear(d_model, d_model)        def forward(self, q, k, v, mask=None):        # 输入维度：q, k, v 均为 [batch_size, seq_len, d_model]        B, T_q, _ = q.shape        T_k = k.shape[1]                # 对 Q、K、V 做线性映射（注意：映射后维度仍是 d_model）        q = self.q_proj(q)        k = self.k_proj(k)        v = self.v_proj(v)                # 拆分多头：[B, T, D] → [B, num_heads, T, head_dim]        def split_heads(x):            return x.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)                q = split_heads(q)        k = split_heads(k)        v = split_heads(v)                # 缩放点积注意力        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)        # 如果提供了 mask，则屏蔽无效位置（如 <pad> 或未来信息）        if mask is not None:            scores = scores.masked_fill(mask == 0, -1e9)                # 计算注意力权重（概率分布）        attn = torch.softmax(scores, dim=-1)                # 使用注意力权重加权值向量，获得每个位置的上下文表示        out = torch.matmul(attn, v)                # 合并多头：[B, num_heads, T, head_dim] → [B, T, D]        out = out.transpose(1, 2).contiguous().view(B, T_q, -1)                # 通过输出线性层整合信息，维度仍为 d_model        return self.out_proj(out)

（3）前馈神经网络与 Transformer Block

在多头注意力之后，Transformer 还会对每个位置的表示单独进行一次非线性变换，这就是所谓的前馈神经网络（Feed-Forward Network，简称 FFN）。

这个模块其实非常简单：对每个位置上的向量先通过一个线性层映射到更高维度，激活之后再映射回来，增强模型的表达能力。

复制代码

class FeedForward(nn.Module):    def __init__(self, d_model, d_ff):        super().__init__()        self.net = nn.Sequential(            nn.Linear(d_model, d_ff),            nn.ReLU(),            nn.Linear(d_ff, d_model),        )        def forward(self, x):        return self.net(x)
class TransformerBlock(nn.Module):    def __init__(self, d_model, num_heads, d_ff):        super().__init__()        # 多头注意力        self.attn = MultiHeadAttention(d_model, num_heads)        # 前馈网络        self.ff = FeedForward(d_model, d_ff)        # 第一次 LayerNorm        self.norm1 = nn.LayerNorm(d_model)        # 第二次 LayerNorm        self.norm2 = nn.LayerNorm(d_model)        def forward(self, x, mask=None):        # 残差连接 + LayerNorm（注意力）        x = self.norm1(x + self.attn(x, x, x, mask))        # 残差连接 + LayerNorm（前馈）        x = self.norm2(x + self.ff(x))        return x

（4）构建编码器和解码器

Transformer 模型由编码器（Encoder） 和解码器（Decoder） 组成。编码器处理源语言输入，提取上下文特征；解码器则根据编码器输出和之前生成的目标词语，逐步生成翻译结果。

编码器模块

编码器由多个 TransformerBlock 组成，每个 block 使用共享的参数结构。

复制代码

class Encoder(nn.Module):    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len=100):        super().__init__()        # 词嵌入层：将输入的词ID转为向量表示        self.embedding = nn.Embedding(vocab_size, d_model)        # 位置编码：添加序列中单词的位置信息        self.pos_enc = PositionalEncoding(d_model, max_len)        # 堆叠多个 TransformerBlock（多层编码器）        self.layers = nn.ModuleList([            TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers)        ])        def forward(self, x, mask=None):        # 输入 x: [batch_size, seq_len] -> 词ID序列        x = self.embedding(x)        x = self.pos_enc(x)        # 依次通过每一层 TransformerBlock        for layer in self.layers:            x = layer(x, mask)        return x

（2）解码器模块（支持掩码）

解码器与编码器结构类似，但每个 TransformerBlock 包括两个注意力子层：

Masked Self-Attention：防止模型"看到"未来词语；
Encoder-Decoder Attention：使解码器能关注输入序列的编码结果。

class DecoderBlock(nn.Module): def init(self, d_model, num_heads, d_ff): super().init() # 解码器中的三个子层：Masked Self-Attention、Encoder-Decoder Attention、前馈网络 self.self_attn = MultiHeadAttention(d_model, num_heads) self.enc_dec_attn = MultiHeadAttention(d_model, num_heads) self.ff = FeedForward(d_model, d_ff) # 每个子层后都跟 LayerNorm 和残差连接 self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) def forward(self, x, enc_out, tgt_mask=None, src_mask=None): # Masked Self-Attention：仅关注当前位置及之前的 token x = self.norm1(x + self.self_attn(x, x, x, mask=tgt_mask)) # Encoder-Decoder Attention：让 decoder 看 encoder 的输出 x = self.norm2(x + self.enc_dec_attn(x, enc_out, enc_out, mask=src_mask)) # 前馈网络 + 残差连接 + LayerNorm x = self.norm3(x + self.ff(x)) return x
class Decoder(nn.Module): def init(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_len=100): super().init() # 解码器的词嵌入层 self.embedding = nn.Embedding(vocab_size, d_model) # 位置编码，表示目标序列的位置信息 self.pos_enc = PositionalEncoding(d_model, max_len) # 堆叠多个 DecoderBlock（多层解码器） self.layers = nn.ModuleList([ DecoderBlock(d_model, num_heads, d_ff) for _ in range(num_layers) ]) # 输出层，将 decoder 输出映射为词表大小，用于 softmax 预测词 self.fc_out = nn.Linear(d_model, vocab_size) def forward(self, x, enc_out, tgt_mask=None, src_mask=None): x = self.embedding(x) x = self.pos_enc(x) # 依次通过每一层 DecoderBlock for layer in self.layers: x = layer(x, enc_out, tgt_mask, src_mask) # 最后输出分类得分，用于预测词 return self.fc_out(x)

（5）整合完整的 Transformer 模型

我们将编码器和解码器组合为一个完整的 Transformer 模型。

复制代码

class Transformer(nn.Module):    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=128, num_heads=4, d_ff=512, num_layers=2, max_len=100):        super().__init__()        self.encoder = Encoder(src_vocab_size, d_model, num_heads, d_ff, num_layers, max_len)        self.decoder = Decoder(tgt_vocab_size, d_model, num_heads, d_ff, num_layers, max_len)        def forward(self, src, tgt, src_mask=None, tgt_mask=None):        enc_out = self.encoder(src, src_mask)        out = self.decoder(tgt, enc_out, tgt_mask, src_mask)        return out

三、模型训练与掩码机制

在训练 Transformer 模型时，为了让模型正确学习序列依赖关系，我们需要构建合适的掩码（mask），分别用于：

源语言掩码 ：屏蔽掉<pad>填充位，防止编码器在注意力中处理无效内容；
目标语言掩码 ：构造下三角掩码，确保模型在预测每个词时只能访问该词之前的内容，防止信息泄露。

（一）构建掩码函数

复制代码

# 构造源语言的掩码：True 表示有效位置，False 表示 <pad>def create_src_mask(src_tensor, pad_idx):    mask = (src_tensor != pad_idx).unsqueeze(1).unsqueeze(2)    return mask
# 构造目标语言的下三角掩码：防止看到未来词def generate_subsequent_mask(size):    return torch.tril(torch.ones(size, size)).unsqueeze(0).unsqueeze(0)

下三角掩码本质上是一个下三角矩阵（lower triangular matrix），对角线以下为 1，其余为 0。

举例：若目标序列长度为 5，则掩码为：

复制代码

[[1, 0, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 1, 0, 0], [1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]

这个矩阵的作用是：

第 1 个词只能看到自己；
第 2 个词可以看到前 2 个；
第 3 个词可以看到前 3 个；
......以此类推。

（二）执行训练循环

我们使用**交叉熵损失函数（CrossEntropyLoss）**来训练翻译模型，其中 <pad> 会被忽略，以防止模型学习无效的填充部分：

复制代码

import torchimport torch.nn as nnimport torch.optim as optim
# 使用 GPU（如果可用）device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 实例化模型：指定源语言和目标语言的词表大小model = Transformer(    src_vocab_size=len(src_vocab),    tgt_vocab_size=len(tgt_vocab),    d_model=128,    num_heads=4,    d_ff=512,    num_layers=2).to(device)
# 优化器：Adam 通常在 Transformer 中表现良好optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 损失函数：忽略目标序列中的 <pad> 位置loss_fn = nn.CrossEntropyLoss(ignore_index=tgt_vocab["<pad>"])
num_epochs = 100for epoch in range(num_epochs):    model.train()    total_loss = 0    # 遍历每个样本（因数据量小，这里不使用批量训练）    for src_seq, tgt_seq in zip(input_seqs, target_seqs):        # 转换为张量并移动到设备        src_tensor = torch.tensor([src_seq], dtype=torch.long).to(device)        tgt_tensor = torch.tensor([tgt_seq], dtype=torch.long).to(device)                # 构造 decoder 的输入和目标        # 输入：去掉句末 <eos>        # 输出：去掉句首 <bos>        tgt_input = tgt_tensor[:, :-1]        tgt_output = tgt_tensor[:, 1:]
        # 生成源语言的注意力掩码（mask <pad>），防止编码器注意力聚焦到无效位置        src_mask = create_src_mask(src_tensor, src_vocab["<pad>"]).to(device)        # 生成 mask：解码器自注意力中的下三角 mask（防止看到未来词）        tgt_mask = generate_subsequent_mask(tgt_input.size(1)).to(device)
        # 前向传播        logits = model(src_tensor, tgt_input, src_mask=src_mask, tgt_mask=tgt_mask)        # 将输出 reshape 为 [batch*seq_len, vocab_size]，与标签对齐        logits = logits.reshape(-1, logits.size(-1))        tgt_output = tgt_output.reshape(-1)                # 计算损失        loss = loss_fn(logits, tgt_output)                # 反向传播与优化        optimizer.zero_grad()        loss.backward()        optimizer.step()                # 累计损失        total_loss += loss.item()        if (epoch + 1) % 10 == 0 or epoch == 0:        print(f"Epoch {epoch+1:>3}: Loss = {total_loss:.4f}")

五、推理解码：贪婪搜索

模型训练完成后，我们希望能够输入一句英文，得到对应的中文翻译。这就需要构建一个推理过程。

我们采用最简单的**贪婪解码（Greedy Decoding）**策略：每一步选择当前概率最高的词，逐步生成完整句子，直到遇到 <eos> 或达到最大长度。

复制代码

def greedy_decode(model, input_sentence, max_len=10):    model.eval()    # 将输入英文句子编码为 token ID 序列，并转为张量，形状：[1, seq_len]    input_ids = torch.tensor([encode(input_sentence, src_vocab)], dtype=torch.long).to(device)        # 在推理阶段关闭梯度计算，加快速度、节省显存    with torch.no_grad():        # 生成源语言的注意力掩码（mask <pad>），防止编码器注意力聚焦到无效位置        src_mask = create_src_mask(input_ids, src_vocab["<pad>"]).to(device)        # 1. 编码器部分：将输入序列传入编码器，得到上下文表示 enc_out        enc_out = model.encoder(input_ids, mask=src_mask)        # 初始化生成序列，起始 token 为 <bos>（句子开头）        decoded = [tgt_vocab["<bos>"]]                # 2. 解码器部分：逐步生成每一个目标词（最多 max_len 个）        for _ in range(max_len):            # 当前已生成的 token ID 序列作为 decoder 的输入            tgt_input = torch.tensor([decoded], dtype=torch.long).to(device)                        # 为 decoder 自注意力生成下三角 mask（防止看到未来词）            tgt_mask = generate_subsequent_mask(tgt_input.size(1)).to(device)                        # 将当前 decoder 输入和 encoder 输出传入 decoder，得到 logits            logits = model.decoder(tgt_input, enc_out, tgt_mask=tgt_mask)                        # 获取当前时间步（最后一个 token）对应的预测概率最大值            next_token = logits[0, -1].argmax().item()                        # 将预测 token 加入到已生成的序列中            decoded.append(next_token)                        # 如果预测为 <eos>，表示句子结束，停止解码            if next_token == tgt_vocab["<eos>"]:                break        # 构建目标词表的反向映射：ID → 词    inv_vocab = {v: k for k, v in tgt_vocab.items()}        # 将生成的 token ID（去掉 <bos> 和 <eos>）转为词字符串并拼接    return " ".join([inv_vocab.get(tok, "<unk>") for tok in decoded[1:-1]])

示例：英译中翻译效果

我们用训练过的模型尝试翻译一两个句子看看效果（即使是极简模型，也能感受到它的结构能力）：

复制代码

test_sentence = "I love machine learning"output = greedy_decode(model, test_sentence)print("Input:", test_sentence)print("Output:", output)

输出结果：

复制代码

Epoch   1: Loss = 12.8265Epoch  10: Loss = 0.2459Epoch  20: Loss = 0.0639Epoch  30: Loss = 0.0415Epoch  40: Loss = 0.0304Epoch  50: Loss = 0.0237Epoch  60: Loss = 0.0191Epoch  70: Loss = 0.0158Epoch  80: Loss = 0.0134Epoch  90: Loss = 0.0114Epoch 100: Loss = 0.0099Input : I love machine learningOutput: 我 喜欢 机器 学习

当然，由于数据量小、模型简化，我们的目标不是达到实际应用水平，而是完整跑通 Transformer 推理链路，理解其生成机制。

六、总结与未来方向

本文从零构建了一个简化版的 Transformer 英汉翻译系统，完整展示了从数据准备、模型搭建到训练与推理的关键流程。我们手动构造中英文句对，生成词表并实现编码器，将句子转换为固定长度的 token 序列；随后基于 PyTorch 实现了位置编码、多头注意力、前馈网络、编码器与解码器模块，最终组合为完整的 Transformer 模型。

尽管数据量和模型规模有限，本项目成功演示了 Transformer 在机器翻译任务中的基本原理和端到端流程。通过训练与贪婪解码，我们验证了模型的翻译能力，为进一步理解和实践更复杂的 Transformer 应用打下了基础。