手撕Transformer：PyTorch实现与训练实战指南

手撕 Transformer 这件事，很多人一开始会觉得特别难，尤其是看到 QKV、多头注意力、位置编码这些名词堆在一起时，很容易劝退。但等你真正把一个可运行的 PyTorch 代码从头写到尾，再看到 loss 一点点降下去，会发现 Transformer 的骨架其实非常清晰。本文就按照“概念 → 原理 → 代码 → 训练 → 排错”的顺序，带你把 Transformer 从 Token 到训练完整过一遍，所有代码都能直接跑，重点部分我会拆开揉碎讲清楚。

如果你正准备入门 NLP、做文本分类、做序列标注，或者想搞懂 GPT 类的生成模型底层到底在干什么，这篇文章都非常适合。你不需要有很深的深度学习基础，只要会 Python，装好了 PyTorch，能跟着敲代码，就足够了。

1. 背景与核心概念：为什么最后是 Transformer

1.1 从 RNN 到 Transformer：解决的是什么问题

在 Transformer 出现之前，处理序列数据的标配是 RNN（循环神经网络）和它的变体 LSTM、GRU。RNN 的核心思路是按时间步一个个处理输入，当前时刻的隐藏状态依赖上一时刻的输出。

这个结构有一个天然问题：串行计算。第 t 个词必须等第 t-1 个词算完才能开始，很难并行。虽然 LSTM 通过门控机制缓解了长距离依赖问题，但本质上仍然是把信息一步步往后传，序列一长，前面的信息容易衰减，甚至梯度消失。

CNN 也可以用来处理序列，比如 TextCNN，它通过不同尺寸的卷积核提取 n-gram 特征，可以并行计算，但感受野有限，很难直接捕捉很长的依赖关系。

Transformer 的思路完全不同。它抛弃了“按顺序读”的方式，一次性看到整个序列，通过一个叫“自注意力”的机制，直接计算任意两个位置之间的关联权重。也就是说，无论两个词隔得多远，Transformer 都能在一步之内建模它们的关系。同时，由于每个位置的向量都可以独立计算注意力，Transformer 的训练过程可以高度并行，这也是它能在超大语料上训练的基础。

一句话总结：Transformer 用自注意力替代了循环结构，用并行计算替代了串行递推，从而在长序列建模和训练效率上全面超越了 RNN 系列。

1.2 Token 是什么：模型根本读不懂单词

很多刚接触 NLP 的朋友会对 Token 这个概念很模糊。Token 是模型处理文本的最小单位，可以理解为一个“词元”。英文中一个单词通常就是一个 Token，比如hello；中文则可能一个字、一个词或者一个子词是一个 Token，取决于分词方式。

模型是纯数学计算，它读不了字符串。我们必须把文本先切成 Token，再把每个 Token 映射成一个固定维度的向量，这个过程叫 Token Embedding（词嵌入）。在 Transformer 中，Token 化通常是第一步：原始文本 → Token 序列 → Token ID 序列 → 词嵌入矩阵。

举个简单例子，一句话我喜欢学习 Transformer，分词后可能是['我', '喜欢', '学习', 'Transformer']，每个 Token 对应一个 ID，比如[12, 33, 58, 1024]。然后查词表得到一个向量矩阵，形状是[序列长度, 嵌入维度]，这个矩阵就是模型真正能处理的输入。

需要强调一点：Token 的切分方式直接影响模型效果。现在主流是 BPE（字节对编码）或 WordPiece 这种子词切分方式，可以有效处理未登录词。不过本文手写 Demo 为了方便，会直接用字符级或者简单分词，重点在于体会流程。

1.3 Transformer 宏观架构速览

标准的 Transformer 模型（论文《Attention Is All You Need》中的结构）分成两部分：Encoder（编码器）和 Decoder（解码器）。

Encoder 负责把输入序列编码成一组上下文相关的向量表示。它的内部由多个相同的层堆叠而成，每一层包含两个子层：多头自注意力机制（Multi-Head Self-Attention）和前馈神经网络（Feed-Forward Network，FFN）。每个子层外面都套了残差连接和层归一化。

Decoder 负责根据 Encoder 的输出和已经解码出的内容，一步步生成目标序列。它的每一层比 Encoder 多了一个 Masked Multi-Head Attention，也就是掩码自注意力，用来确保预测当前位置时看不到未来位置的信息。

如果你只想做文本分类、语义匹配这类理解任务，只需要 Encoder 部分就够了。如果你想做机器翻译或者类似 GPT 的文本生成，才需要完整的 Encoder-Decoder 结构或只有 Decoder 的架构。本文的实战案例会从 Encoder 入手，因为它更容易理解，也最适合作为手撕 Transformer 的第一步。

2. 环境准备与版本说明

2.1 PyTorch 环境安装

在开始写代码之前，先确保你的电脑上有 Python 和 PyTorch。建议使用 conda 创建独立环境，避免把系统 Python 搞乱。

下面以最常见的 conda 创建方式为例：

conda create -n transformer python=3.10 conda activate transformer pip install torch --index-url https://download.pytorch.org/whl/cpu

如果你有 NVIDIA 显卡，并且已经安装好 CUDA 驱动，可以按官方命令安装 GPU 版本：

pip install torch --index-url https://download.pytorch.org/whl/cu118

版本需要根据你的项目实际情况调整。比如 PyTorch 官方现在已经有很多发行版本，CPU 版和 CUDA 版命令不同。安装完成后，用下面这段代码验证：

import torch print(torch.__version__) print(torch.cuda.is_available())

可以看到类似这样的输出：

2.4.0 True

如果你没有 GPU，完全不影响本文的学习。本文的 Demo 用 CPU 也能在几分钟内跑完。

这里多提一句，很多人在安装 PyTorch 时遇到下载慢的问题，尤其是用默认 PyPI 源的时候。建议换成国内镜像：

pip install torch --index-url https://download.pytorch.org/whl/cpu -i https://pypi.tuna.tsinghua.edu.cn/simple

不过用官方源和镜像源同时指定时需要注意兼容性，如果遇到问题，最简单的方式是直接用官方命令安装。环境问题下面第六章会专门排查。

2.2 需要的 Python 库

本文只需要三个库：

• torch：深度学习框架
• numpy：数值计算，PyTorch 的依赖之一，一般会随 torch 自动装好
• matplotlib：用于可视化 loss 曲线

安装：

pip install matplotlib

2.3 验证最小可运行代码

在正式开始前，写一个最小的 PyTorch 张量检查脚本，确认环境没问题：

import torch x = torch.randn(4, 8) print(x.shape) print(x)

如果能正常输出矩阵，说明环境没问题。接下来我们进入核心原理部分。

3. 核心原理拆解：从 Token 到自注意力

这一节是整个 Transformer 最核心的部分。我会按照数据流的顺序，一步步说明每个模块要做什么、为什么这么做，并给出对应的 PyTorch 代码片段。这些片段最后会组装成一个完整的可训练模型。

3.1 Token Embedding 与输入表示

输入是一组 Token ID，形状是[batch_size, seq_len]。比如一个 batch 有 2 句话，每句话有 5 个 Token，那形状就是[2, 5]。

我们需要一个词嵌入矩阵，把每个 Token ID 映射成一个向量。比如嵌入维度是 128，那词表大小是 10000 的话，嵌入矩阵就是[10000, 128]。

这层在 PyTorch 里直接用nn.Embedding实现：

import torch.nn as nn embedding = nn.Embedding(num_embeddings=10000, embedding_dim=128) x = torch.randint(0, 10000, (2, 5)) out = embedding(x) print(out.shape) # torch.Size([2, 5, 128])

得到的结果out就是每个 Token 的向量表示，形状是[2, 5, 128]，对应[batch_size, seq_len, d_model]。

3.2 位置编码：为什么需要它

自注意力机制有一个特点：它对位置不敏感。也就是说，把句子里的 Token 顺序打乱，注意力计算出来的结果是一样的，因为注意力只关心“Token 之间的相似度”，不关心“谁在前谁在后”。

这显然不行。我喜欢你和你喜欢我的 Token 完全相同，但语义差别很大。RNN 天然有顺序信息，因为它是按时间步输入的，但 Transformer 没有这个特性，所以必须把位置信息“注入”到输入向量中。

最经典的做法是正余弦位置编码（Sinusoidal Positional Encoding）。它的公式如下：

对于位置pos和维度索引i：

• 偶数维：PE(pos, 2i) = sin(pos / 10000^(2i / d_model))
• 奇数维：PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

用代码实现：

import torch import math def sinusoidal_position_encoding(seq_len, d_model): pe = torch.zeros(seq_len, d_model) position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # [1, seq_len, d_model] return pe pe = sinusoidal_position_encoding(10, 128) print(pe.shape) # torch.Size([1, 10, 128])

为什么正余弦编码有效？因为正余弦函数具有周期性，不同位置会产生不同的编码向量，而且模型可以通过线性变换学习到相对位置关系。简单理解：位置编码给每个位置一个独一无二的“坐标”，模型可以通过这个坐标感知词与词之间的距离。

在实际代码中，位置编码是直接加到词嵌入上的：

x = embedding(tokens) # [batch, seq_len, d_model] positions = sinusoidal_position_encoding(seq_len, d_model) x = x + positions

3.3 自注意力机制与 QKV

自注意力是整个 Transformer 的灵魂。它的目的是让序列中的每个 Token 都能“关注”到其他 Token，并根据关注程度聚合信息。

先解释 QKV 这三个字母的含义：

• Q（Query）：查询向量，表示当前 Token 想要从其他 Token 那里获取什么信息。
• K（Key）：键向量，表示当前 Token 能够提供什么信息。
• V（Value）：值向量，表示当前 Token 携带的实际内容。

可以类比搜索引擎：你输入搜索词（Query），系统匹配网页标题（Key），然后返回网页内容（Value）。

计算过程分三步：

第一步，通过三个线性变换，把输入向量x映射成 Q、K、V：

Q = x @ W_q # [batch, seq_len, d_k] K = x @ W_k # [batch, seq_len, d_k] V = x @ W_v # [batch, seq_len, d_v]

其中d_k和d_v是 Q/K 和 V 的维度，通常是嵌入维度除以头数。

第二步，计算注意力分数。Q 和 K 做点积，得到每个 Token 对每个 Token 的原始相关性得分：

scores = Q @ K.transpose(-2, -1) / sqrt(d_k)

除以sqrt(d_k)是为了防止点积结果过大，导致 softmax 梯度消失。这就是论文中提到的 Scaled Dot-Product Attention。

第三步，对 scores 做 softmax 归一化，得到注意力权重，再与 V 加权求和：

weights = softmax(scores, dim=-1) output = weights @ V

整体代码实现：

import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, mask=None): d_k = Q.size(-1) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) weights = F.softmax(scores, dim=-1) output = torch.matmul(weights, V) return output

这里mask用于屏蔽某些位置，比如 Decoder 中防止看到未来位置，或者 Padding 位置不参与注意力计算。mask 为 0 的位置分数会被设置成一个极大的负数，softmax 之后对应概率趋近于 0。

3.4 多头自注意力机制

多头自注意力机制并不是做一次注意力，而是把 Q、K、V 拆成多份，每一份单独做注意力计算，最后拼接起来。

为什么需要多头？因为一次注意力只能学到一种“关系模式”。比如在句子小明喜欢吃苹果，因为苹果很甜中，小明和苹果有关系，喜欢和吃也有关系。多头允许模型在不同子空间里学习不同的关系模式：一个头可能关注语法依赖，另一个头可能关注语义相似性。

PyTorch 实现多头注意力的常用方式是：使用一个大的线性层同时生成所有头的 Q、K、V，然后通过 reshape 拆分。

class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): batch_size, seq_len, _ = x.size() # 生成 Q、K、V 并拆分成多头 Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) weights = F.softmax(scores, dim=-1) weights = self.dropout(weights) context = torch.matmul(weights, V) # [batch, n_heads, seq_len, d_k] context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output = self.W_o(context) return output

拆分后的 Q/K/V 维度是[batch_size, n_heads, seq_len, d_k]。注意力计算在最后两个维度上进行，也就是每个头独立计算。最后把多个头的输出拼回[batch_size, seq_len, d_model]，再经过一个输出线性变换。

3.5 层归一化与残差连接

Transformer 的每一层都不是单独的一个模块，而是“残差连接 + 层归一化 + 子层”的组合。

残差连接，也就是output = x + sublayer(x)。它解决的是深层网络训练时梯度消失的问题。如果子层学习到的变换接近恒等映射，残差连接可以保证梯度顺畅地回传。

层归一化（Layer Normalization，简称 LayerNorm）对每个样本的每个位置单独做归一化，即在最后一维上计算均值和方差。这比 BatchNorm 更稳定，因为不依赖 batch 大小，也不会受序列长度不一致影响。

class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.attention = MultiHeadAttention(d_model, n_heads, dropout) self.norm1 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): # 第一个子层：多头注意力 attn_output = self.attention(x, mask) x = self.norm1(x + self.dropout(attn_output)) # 第二个子层：前馈网络 ffn_output = self.ffn(x) x = self.norm2(x + self.dropout(ffn_output)) return x

注意这里的写法是 Post-Norm，也就是先残差再归一化。论文原文用的是这种方式。实际工程中也有 Pre-Norm（先归一化再子层），它在深层网络中更稳定。本文用 Post-Norm 是为了贴近原论文。

3.6 前馈网络 FFN

FFN 是一个两层的全连接网络，中间接一个 ReLU 激活函数。公式如下：

FFN(x) = max(0, xW1 + b1)W2 + b2

它的作用是对注意力聚合后的信息做非线性变换，增强模型的表达能力。每个位置的 FFN 参数共享，但不同位置独立计算。

self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), )

d_ff一般是d_model的 4 倍，比如 d_model=128 时，d_ff=512。

4. 完整实战：手写一个可训练的 Transformer 分类模型

这一节我们把前面拆开的模块组装起来，做一个文本分类 Demo。数据集用一个小规模的英文情感分类数据，包含正面和负面句子。模型结构是 Encoder-only 的 Transformer，最后接一个全连接分类头。

4.1 整体代码结构与流程

整个项目只有一个 Python 文件，建议命名为transformer_demo.py。流程如下：

1. 构建一个小型数据集
2. 文本 Token 化，并转换成 Token ID
3. 定义 Transformer 模型
4. 定义损失函数、优化器
5. 训练并输出 loss 和准确率
6. 验证模型效果

4.2 构造演示数据

为了可运行性和可复现性，我们自己构造 100 条英文短句，模拟正面和负面情感。这里使用最简单的空格分词。

import random import torch import torch.nn as nn import torch.nn.functional as F import math # 固定随机种子，方便复现 random.seed(42) torch.manual_seed(42) positive_sentences = [ "i love this movie", "the film is great", "what a wonderful experience", "i am so happy", "this is amazing", "the service was excellent", "best day ever", "i really enjoyed it", "the food is delicious", "highly recommended", ] negative_sentences = [ "i hate this movie", "the film is terrible", "what a bad experience", "i am so sad", "this is awful", "the service was poor", "worst day ever", "i really disliked it", "the food is awful", "not recommended", ] data = [] for sent in positive_sentences: data.append((sent, 1)) for sent in negative_sentences: data.append((sent, 0)) # 扩展到 100 条 all_data = [] for _ in range(5): for item in data: all_data.append(item)

4.3 构建词表与 Token 化

为了简单，我们直接给每个单词分配一个 ID，并添加<pad>和<unk>两个特殊 Token。

def build_vocab(data): vocab = {"<pad>": 0, "<unk>": 1} for sent, _ in data: for word in sent.split(): if word not in vocab: vocab[word] = len(vocab) return vocab def encode_sentence(sent, vocab, max_len): tokens = [vocab.get(word, vocab["<unk>"]) for word in sent.split()] if len(tokens) > max_len: tokens = tokens[:max_len] else: tokens = tokens + [vocab["<pad>"]] * (max_len - len(tokens)) return torch.tensor(tokens, dtype=torch.long)

这里max_len取数据集中最长句子的长度。为了避免 padding 过多，我们统一设置max_len=8。

4.4 定义 Transformer 模型

把我们之前写的三个类整合到一起：MultiHeadAttention、TransformerBlock、TransformerEncoder。最后加上分类头。

class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len, num_classes, dropout=0.1): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.positional_encoding = sinusoidal_position_encoding(max_len, d_model) self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers) ]) self.norm = nn.LayerNorm(d_model) self.fc = nn.Linear(d_model, num_classes) def forward(self, x, mask=None): # x: [batch, seq_len] embedded = self.embedding(x) # [batch, seq_len, d_model] x = embedded + self.positional_encoding[:, :x.size(1), :].to(x.device) for block in self.blocks: x = block(x, mask) x = self.norm(x) # 取序列第一个位置的表示作为整体语义表示 pooled = x[:, 0, :] output = self.fc(pooled) return output

关于取每个样本的哪个 Token 做分类，常见的做法有两种：一种是取第一个 Token（类似 BERT 的[CLS]），另一种是对所有 Token 做平均池化。这里简单起见，取第一个 Token。

如果你要做文本生成，只需要把self.fc换成一个映射到词表的线性层，并调整训练目标为交叉熵。但本文先讲分类，生成留到扩展部分。

4.5 准备训练函数

def train_model(model, data, vocab, max_len, epochs=20, batch_size=16, lr=0.001): optimizer = torch.optim.Adam(model.parameters(), lr=lr) loss_fn = nn.CrossEntropyLoss() dataset_size = len(data) losses = [] for epoch in range(epochs): # 每个 epoch 打乱数据 random.shuffle(data) total_loss = 0.0 correct = 0 total = 0 for i in range(0, dataset_size, batch_size): batch = data[i:i+batch_size] inputs = torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels = torch.tensor([label for _, label in batch], dtype=torch.long) optimizer.zero_grad() outputs = model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * len(batch) preds = outputs.argmax(dim=-1) correct += (preds == labels).sum().item() total += len(batch) avg_loss = total_loss / total acc = correct / total losses.append(avg_loss) print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}") return losses

4.6 跑起来

vocab = build_vocab(all_data) model = TransformerEncoder( vocab_size=len(vocab), d_model=64, n_heads=4, d_ff=256, n_layers=2, max_len=8, num_classes=2, dropout=0.1, ) losses = train_model(model, all_data, vocab, max_len=8, epochs=20)

模型参数量非常小，CPU 上训练也很快。预期 loss 会逐步下降，准确率接近 100%，因为数据是我们构造的，任务本身比较简单。

4.7 结果说明

训练结束后，你可以看到类似这样的输出：

Epoch 1/20 - Loss: 0.6931 - Acc: 0.5100 Epoch 5/20 - Loss: 0.5972 - Acc: 0.7200 Epoch 10/20 - Loss: 0.3871 - Acc: 0.9200 Epoch 15/20 - Loss: 0.1682 - Acc: 0.9800 Epoch 20/20 - Loss: 0.0577 - Acc: 1.0000

可以看到模型能够学习到这个简单的分类任务。这说明 Transformer 模型的核心组件（自注意力、位置编码、FFN）已经正确组装并能通过反向传播更新参数。

如果准确率一直上不去，要么是学习率设置不当，要么是模型结构有 bug，这个问题在第六章会详细排查。

4.8 完整代码汇总

为了方便复制，我把完整代码整合在下面。你只需要把这段代码保存为transformer_demo.py，直接python transformer_demo.py即可运行。

import random import math import torch import torch.nn as nn import torch.nn.functional as F # ---------- 固定随机种子 ---------- random.seed(42) torch.manual_seed(42) # ---------- 1. 构造数据 ---------- positive_sentences = [ "i love this movie", "the film is great", "what a wonderful experience", "i am so happy", "this is amazing", "the service was excellent", "best day ever", "i really enjoyed it", "the food is delicious", "highly recommended", ] negative_sentences = [ "i hate this movie", "the film is terrible", "what a bad experience", "i am so sad", "this is awful", "the service was poor", "worst day ever", "i really disliked it", "the food is awful", "not recommended", ] data = [] for sent in positive_sentences: data.append((sent, 1)) for sent in negative_sentences: data.append((sent, 0)) all_data = [] for _ in range(5): for item in data: all_data.append(item) # ---------- 2. 构建词表和编码 ---------- def build_vocab(data): vocab = {"<pad>": 0, "<unk>": 1} for sent, _ in data: for word in sent.split(): if word not in vocab: vocab[word] = len(vocab) return vocab def encode_sentence(sent, vocab, max_len): tokens = [vocab.get(word, vocab["<unk>"]) for word in sent.split()] if len(tokens) > max_len: tokens = tokens[:max_len] else: tokens = tokens + [vocab["<pad>"]] * (max_len - len(tokens)) return torch.tensor(tokens, dtype=torch.long) # ---------- 3. 位置编码 ---------- def sinusoidal_position_encoding(seq_len, d_model): pe = torch.zeros(seq_len, d_model) position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) return pe # ---------- 4. 多头注意力 ---------- class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): batch_size, seq_len, _ = x.size() Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) weights = F.softmax(scores, dim=-1) weights = self.dropout(weights) context = torch.matmul(weights, V) context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output = self.W_o(context) return output # ---------- 5. Transformer 编码块 ---------- class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.attention = MultiHeadAttention(d_model, n_heads, dropout) self.norm1 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): attn_output = self.attention(x, mask) x = self.norm1(x + self.dropout(attn_output)) ffn_output = self.ffn(x) x = self.norm2(x + self.dropout(ffn_output)) return x # ---------- 6. Transformer Encoder ---------- class TransformerEncoder(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len, num_classes, dropout=0.1): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.positional_encoding = sinusoidal_position_encoding(max_len, d_model) self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers) ]) self.norm = nn.LayerNorm(d_model) self.fc = nn.Linear(d_model, num_classes) def forward(self, x, mask=None): embedded = self.embedding(x) x = embedded + self.positional_encoding[:, :x.size(1), :].to(x.device) for block in self.blocks: x = block(x, mask) x = self.norm(x) pooled = x[:, 0, :] output = self.fc(pooled) return output # ---------- 7. 训练 ---------- def train_model(model, data, vocab, max_len, epochs=20, batch_size=16, lr=0.001): optimizer = torch.optim.Adam(model.parameters(), lr=lr) loss_fn = nn.CrossEntropyLoss() dataset_size = len(data) for epoch in range(epochs): random.shuffle(data) total_loss = 0.0 correct = 0 total = 0 for i in range(0, dataset_size, batch_size): batch = data[i:i+batch_size] inputs = torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels = torch.tensor([label for _, label in batch], dtype=torch.long) optimizer.zero_grad() outputs = model(inputs) loss = loss_fn(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * len(batch) preds = outputs.argmax(dim=-1) correct += (preds == labels).sum().item() total += len(batch) avg_loss = total_loss / total acc = correct / total print(f"Epoch {epoch+1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}") # ---------- 8. 启动训练 ---------- vocab = build_vocab(all_data) print("词表大小:", len(vocab)) model = TransformerEncoder( vocab_size=len(vocab), d_model=64, n_heads=4, d_ff=256, n_layers=2, max_len=8, num_classes=2, dropout=0.1, ) print("模型参数总量:", sum(p.numel() for p in model.parameters())) train_model(model, all_data, vocab, max_len=8, epochs=20)

5. 从 Encoder 到 Decoder：生成任务怎么做

很多读者学完 Encoder 之后，会想知道 GPT 那种文本生成模型是怎么做出来的。这里做一个概念性的扩展。

Decoder 和 Encoder 最大的区别有两点：

第一，Decoder 引入了掩码自注意力（Masked Self-Attention）。在预测第 t 个 Token 时，模型只能看到第 1 到第 t-1 个 Token，不能看到未来信息。实现方式是在 softmax 之前，把当前时刻之后位置的注意力分数设为一个非常大的负数。这个操作也叫因果掩码（Causal Mask）。

def generate_causal_mask(seq_len): mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() return ~mask

第二，生成是逐步进行的。训练时我们可以用 Teacher Forcing 一次性把完整目标序列输入模型；但推理时模型只能先预测第一个 Token，把它拼到输入后面，再预测第二个 Token，依次往复。

如果未来你想做机器翻译或者写一个 GPT 风格的小模型，可以在本文 Encoder 代码的基础上，加入因果掩码，把分类头换成词表大小的线性层，然后使用交叉熵损失对下一个 Token 进行预测。本文的代码已经把自注意力、位置编码这些核心模块都实现了，扩展成 Decoder 的难度并不会太大。

6. 常见问题与排查思路

前面代码虽然能跑通，但你在自己写或者调整的时候，大概率还会碰到各种问题。下面列出高频问题。

针对 Q/K/V 维度问题，我给一个排查模板。在MultiHeadAttention.forward中加入下面的打印：

print("Q shape:", Q.shape) print("K shape:", K.shape) print("V shape:", V.shape) print("scores shape:", scores.shape)

正确的输出应该是：

Q shape: torch.Size([16, 4, 8, 16]) K shape: torch.Size([16, 4, 8, 16]) V shape: torch.Size([16, 4, 8, 16]) scores shape: torch.Size([16, 4, 8, 8])

这里的 16 是 batch size，4 是头数，8 是序列长度，16 是d_k = d_model / n_heads = 64 / 4。如果形状不是这样，说明拆头或者 transpose 的逻辑有误。

另一个容易踩的坑是 GPU 和 CPU 的张量设备不一致。当你在TransformerEncoder.forward里做embedded + self.positional_encoding时，如果模型在 GPU 上而positional_encoding在 CPU 上，PyTorch 会直接报错。解决方案是在 forward 中动态转移设备：

x = embedded + self.positional_encoding[:, :x.size(1), :].to(x.device)

7. 最佳实践与工程建议

7.1 代码层面的建议

手撕完一个最小模型后，往工程方向走时要注意几点。

第一，d_model 和 n_heads 的配合。d_model 必须能够被 n_heads 整除，否则无法拆头。常见组合是 d_model=512、n_heads=8，或者 d_model=768、n_heads=12。这个数字不是随便定的，它保证了每个头有足够的信息容量，同时拆分后计算效率较高。

第二，dropout 不要滥用。在训练时，Transformer 通常会在注意力权重、残差输出、FFN 内部都加 dropout。但推理时要切到model.eval()模式，否则 dropout 会保留，导致推理结果不稳定。

model.eval() with torch.no_grad(): outputs = model(inputs)

第三，对于真实项目，建议直接使用nn.TransformerEncoderLayer或 HuggingFace 的 Transformer 库，而不是自己实现。自己实现的目的是理解原理，工程上使用成熟的库是更稳的选择。

7.2 数据层面的建议

Padding 是 NLP 里绕不开的处理。当一个 batch 中句子长度不一致时，短的句子需要补到和最长句子一样长。但 padding Token 不参与语义计算，所以最好在注意力计算时传入src_key_padding_mask，把 padding 位置直接屏蔽掉。

def make_padding_mask(tokens, pad_idx=0): return (tokens != pad_idx).unsqueeze(1).unsqueeze(2)

这个 mask 的形状是[batch_size, 1, 1, seq_len]，可以广播到注意力分数矩阵[batch_size, n_heads, seq_len, seq_len]。mask 为 False 的位置会被替换成极大负数，softmax 之后权重为 0。

7.3 模型训练与超参数调优

Transformer 对超参数比较敏感。以下是个人实践中的经验总结：

• 学习率：默认 1e-4 到 5e-4，配合 Adam 优化器和 warmup 策略（先线性上升再衰减）效果更好。
• Batch size：显存允许的情况下尽量大，一般 32 到 128 之间。分类任务小模型可以用 16。
• 层数：小数据集 2 到 4 层就够了，堆太多层反而容易过拟合。
• 梯度裁剪：训练生成模型时，clip_grad_norm_几乎是必加的，能防止训练崩溃。
• 监控指标：除了 loss，建议每几个 epoch 在验证集上看准确率。如果训练 loss 下降但验证 loss 上升，说明过拟合，可以加大 dropout 或减小模型规模。

7.4 一个更实用的训练模式

如果你想小步验证自己的代码是否正确，可以加一个梯度检查，或者用“先跑一个过拟合小样本”的方法。具体做法是：只取 10 条训练数据，让模型重复训练 50 个 epoch。如果模型能在这 10 条数据上达到 100% 准确率，说明模型结构没问题。如果连小样本都过拟合不了，问题一定出在代码逻辑上。

这个方法在深度学习实战中非常实用。很多时候我们会在大型数据集上等半天，结果发现是模型维度写错了。先用小样本验证，能节省大量时间。

8. 总结与学习路线

通过本文，你应该已经理解了 Transformer 的完整数据流：文本先变成 Token ID，再通过词嵌入变成向量，加上位置编码后进入多头自注意力层，注意力输出经过残差连接和层归一化，再通过前馈网络，最后送到分类头或词表。

我们也用 PyTorch 从零手写了多头注意力、TransformerBlock、位置编码和一个小型 Encoder 分类模型，并且成功训练到 loss 下降、准确率接近 100%。这说明你真正理解了 Transformer 的核心结构，而不只是会调用nn.Transformer。

接下来，你可以从以下几个方向继续深入：

• 把分类模型改成文本生成模型：加入因果掩码，用上一个 Token 预测下一个 Token，这是 GPT 的雏形。
• 阅读《Attention Is All You Need》原文和 The Illustrated Transformer，把位置编码的数学推导、Scaled Dot-Product Attention 的证明看明白。
• 尝试用nn.TransformerEncoderLayer替换自己写的模块，对比训练速度和效果差异。
• 在真实数据集上做情感分类，比如 IMDb 电影评论，体验从数据处理到模型训练和评估的完整流程。
• 学习 BERT 和 GPT 的源码，HuggingFace 的transformers库是很好的参考。

最后想说的是，手撕 Transformer 不是目的，目的是通过亲手写代码理解深度学习模型是怎么一层层组装出来的。等你自己能独立改代码、调参、加 Module 的时候，你会发现 Transformer 并没有想象中那么神秘。

如果本文对你有帮助，可以收藏备用，以后调参或者面试复习的时候翻出来看看。也欢迎在评论区交流你在跑代码过程中遇到的具体问题。