从零手写Transformer_使用PyTorch实现简化版英汉机器翻译
在上一篇文章中，我们手动实现了 Transformer 的核心组件，并详细解析了其结构原理。本文将在此基础上，演示一个简化版的英汉机器翻译系统，完整呈现其训练与推理流程。我们将使用 PyTorch 逐步搭建整个流程，重点突出核心逻辑，代码简洁易懂，便于理解。
需要说明的是，本文的目标并非构建高精度的翻译系统，而是通过最小样本和最简配置，演示 Transformer 从数据准备到模型构建，再到生成翻译结果的完整闭环。
1. 准备数据：构建中英文平行语料
在训练 Transformer 模型之前，我们需要准备一组中英文对照的数据，并实现基本的编码器（Tokenizer），将文本转换为模型可读的数字序列。
1.1 构建中英文句子对
我们先手动构建一个小型中英文平行语料，数据仅用于结构验证，并非用于泛化或真实训练。
raw_data = [
("I love machine learning", "我 喜欢 机器 学习"),
("Deep learning is powerful", "深度 学习 很 强大"),
("Transformer changed everything", "Transformer 改变 了 一切"),
]
这里每个元组表示一组输入输出对，左边是英文句子，右边是中文翻译（已用空格分词，便于处理）。
1.2 构建词表
接下来我们为每种语言分别构建词表（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，以便后续处理。
1.3 编码器（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
1.4 编码中英文数据
将原始句子编码为固定长度的整数序列，作为模型的输入和目标输出：
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 模型的训练。
2. 定义 Transformer 模型结构
我们将使用 PyTorch 搭建完整的 Transformer 模型，包括嵌入层、位置编码、注意力机制、前馈网络、编码器和解码器。
2.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.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)
2.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
2.4 构建编码器和解码器
Transformer 模型由编码器（Encoder）和解码器（Decoder）组成。编码器处理源语言输入，提取上下文特征；解码器则根据编码器输出和之前生成的目标词语，逐步生成翻译结果。
2.4.1 编码器模块
编码器由多个 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.4.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)
2.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
3. 模型训练与掩码机制
在训练 Transformer 模型时，为了让模型正确学习序列依赖关系，我们需要构建合适的掩码（mask），分别用于：
- 源语言掩码：屏蔽掉
<pad>填充位，防止编码器在注意力中处理无效内容； - 目标语言掩码：构造下三角掩码，确保模型在预测每个词时只能访问该词之前的内容，防止信息泄露。
3.1 构建掩码函数
# 构造源语言的掩码：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 个；
- ……以此类推。
3.2 执行训练循环
我们使用**交叉熵损失函数（CrossEntropyLoss）**来训练翻译模型，其中 <pad> 会被忽略，以防止模型学习无效的填充部分：
import torch
import torch.nn as nn
import 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 = 100
for 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}")
4. 推理解码：贪婪搜索
模型训练完成后，我们希望能够输入一句英文，得到对应的中文翻译。这就需要构建一个推理过程。
我们采用最简单的**贪婪解码（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]])
5. 示例：英译中翻译效果
我们用训练过的模型尝试翻译一两个句子看看效果（即使是极简模型，也能感受到它的结构能力）：
test_sentence = "I love machine learning"
output = greedy_decode(model, test_sentence)
print("Input:", test_sentence)
print("Output:", output)
输出结果：
Epoch 1: Loss = 12.8265
Epoch 10: Loss = 0.2459
Epoch 20: Loss = 0.0639
Epoch 30: Loss = 0.0415
Epoch 40: Loss = 0.0304
Epoch 50: Loss = 0.0237
Epoch 60: Loss = 0.0191
Epoch 70: Loss = 0.0158
Epoch 80: Loss = 0.0134
Epoch 90: Loss = 0.0114
Epoch 100: Loss = 0.0099
Input: I love machine learning
Output: 我 喜欢 机器 学习
当然，由于数据量小、模型简化，我们的目标不是达到实际应用水平，而是完整跑通 Transformer 推理链路，理解其生成机制。
6. 总结与未来方向
本文从零构建了一个简化版的 Transformer 英汉翻译系统，完整展示了从数据准备、模型搭建到训练与推理的关键流程。我们手动构造中英文句对，生成词表并实现编码器，将句子转换为固定长度的 token 序列；随后基于 PyTorch 实现了位置编码、多头注意力、前馈网络、编码器与解码器模块，最终组合为完整的 Transformer 模型。
尽管数据量和模型规模有限，本项目成功演示了 Transformer 在机器翻译任务中的基本原理和端到端流程。通过训练与贪婪解码，我们验证了模型的翻译能力，为进一步理解和实践更复杂的 Transformer 应用打下了基础。
写这篇文章花了我很多时间与心力，如果您觉得内容对您有所启发，欢迎点赞并关注「凝思 AI」——您的支持，是我持续创作优质内容的不竭动力！
专注于人工智能与语言、视觉、数据交汇处的前沿探索，涵盖大语言模型（LLMs）、多模态模型、智能体框架、机器学习方法，以及基于数据的潜在空间建模等。