手撕 Transformer 这件事很多人一开始会觉得特别难尤其是看到 QKV、多头注意力、位置编码这些名词堆在一起时很容易劝退。但等你真正把一个可运行的 PyTorch 代码从头写到尾再看到 loss 一点点降下去会发现 Transformer 的骨架其实非常清晰。本文就按照“概念 → 原理 → 代码 → 训练 → 排错”的顺序带你把 Transformer 从 Token 到训练完整过一遍所有代码都能直接跑重点部分我会拆开揉碎讲清楚。如果你正准备入门 NLP、做文本分类、做序列标注或者想搞懂 GPT 类的生成模型底层到底在干什么这篇文章都非常适合。你不需要有很深的深度学习基础只要会 Python装好了 PyTorch能跟着敲代码就足够了。1. 背景与核心概念为什么最后是 Transformer1.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 NetworkFFN。每个子层外面都套了残差连接和层归一化。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 python3.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 matplotlib2.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_embeddings10000, embedding_dim128) 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, 2i1) 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, dtypetorch.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 positions3.3 自注意力机制与 QKV自注意力是整个 Transformer 的灵魂。它的目的是让序列中的每个 Token 都能“关注”到其他 Token并根据关注程度聚合信息。先解释 QKV 这三个字母的含义QQuery查询向量表示当前 Token 想要从其他 Token 那里获取什么信息。KKey键向量表示当前 Token 能够提供什么信息。VValue值向量表示当前 Token 携带的实际内容。可以类比搜索引擎你输入搜索词Query系统匹配网页标题Key然后返回网页内容Value。计算过程分三步第一步通过三个线性变换把输入向量x映射成 Q、K、VQ 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, maskNone): 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, dropout0.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, maskNone): 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, dropout0.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, maskNone): # 第一个子层多头注意力 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 前馈网络 FFNFFN 是一个两层的全连接网络中间接一个 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_model128 时d_ff512。4. 完整实战手写一个可训练的 Transformer 分类模型这一节我们把前面拆开的模块组装起来做一个文本分类 Demo。数据集用一个小规模的英文情感分类数据包含正面和负面句子。模型结构是 Encoder-only 的 Transformer最后接一个全连接分类头。4.1 整体代码结构与流程整个项目只有一个 Python 文件建议命名为transformer_demo.py。流程如下构建一个小型数据集文本 Token 化并转换成 Token ID定义 Transformer 模型定义损失函数、优化器训练并输出 loss 和准确率验证模型效果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, dtypetorch.long)这里max_len取数据集中最长句子的长度。为了避免 padding 过多我们统一设置max_len8。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, dropout0.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, maskNone): # 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, epochs20, batch_size16, lr0.001): optimizer torch.optim.Adam(model.parameters(), lrlr) 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:ibatch_size] inputs torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels torch.tensor([label for _, label in batch], dtypetorch.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(fEpoch {epoch1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}) return losses4.6 跑起来vocab build_vocab(all_data) model TransformerEncoder( vocab_sizelen(vocab), d_model64, n_heads4, d_ff256, n_layers2, max_len8, num_classes2, dropout0.1, ) losses train_model(model, all_data, vocab, max_len8, epochs20)模型参数量非常小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, dtypetorch.long) # ---------- 3. 位置编码 ---------- def sinusoidal_position_encoding(seq_len, d_model): pe torch.zeros(seq_len, d_model) position torch.arange(0, seq_len, dtypetorch.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, dropout0.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, maskNone): 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, dropout0.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, maskNone): 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, dropout0.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, maskNone): 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, epochs20, batch_size16, lr0.001): optimizer torch.optim.Adam(model.parameters(), lrlr) 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:ibatch_size] inputs torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels torch.tensor([label for _, label in batch], dtypetorch.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(fEpoch {epoch1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}) # ---------- 8. 启动训练 ---------- vocab build_vocab(all_data) print(词表大小:, len(vocab)) model TransformerEncoder( vocab_sizelen(vocab), d_model64, n_heads4, d_ff256, n_layers2, max_len8, num_classes2, dropout0.1, ) print(模型参数总量:, sum(p.numel() for p in model.parameters())) train_model(model, all_data, vocab, max_len8, epochs20)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), diagonal1).bool() return ~mask第二生成是逐步进行的。训练时我们可以用 Teacher Forcing 一次性把完整目标序列输入模型但推理时模型只能先预测第一个 Token把它拼到输入后面再预测第二个 Token依次往复。如果未来你想做机器翻译或者写一个 GPT 风格的小模型可以在本文 Encoder 代码的基础上加入因果掩码把分类头换成词表大小的线性层然后使用交叉熵损失对下一个 Token 进行预测。本文的代码已经把自注意力、位置编码这些核心模块都实现了扩展成 Decoder 的难度并不会太大。6. 常见问题与排查思路前面代码虽然能跑通但你在自己写或者调整的时候大概率还会碰到各种问题。下面列出高频问题。问题现象常见原因解决思路loss 不下降学习率过大或过小尝试 1e-4 到 1e-2 之间的学习率观察前几个 step 的 loss 变化loss 为 NaN梯度爆炸或数据含 NaN检查输入是否含 NaN添加梯度裁剪clip_grad_norm_(model.parameters(), 1.0)准确率一直 50%模型结构 bug 或标签不平衡打印模型的 forward 输出形状确认最后一层输出和标签维度匹配训练很慢没有用 GPU 或 batch 太大确认torch.cuda.is_available()或者减小 batch size位置编码加到输入后结果不对位置编码没有转移到 GPU用.to(x.device)同步位置编码和输入张量注意力矩阵维度不匹配Q/K/V 拆头维度写错重点检查view和transpose后的形状建议每步都打印 shape安装 PyTorch 太慢网络原因或源问题用国内镜像或直接去官方站下载本地 wheel 安装针对 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 size4 是头数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_model512、n_heads8或者 d_model768、n_heads12。这个数字不是随便定的它保证了每个头有足够的信息容量同时拆分后计算效率较高。第二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_idx0): 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 并没有想象中那么神秘。如果本文对你有帮助可以收藏备用以后调参或者面试复习的时候翻出来看看。也欢迎在评论区交流你在跑代码过程中遇到的具体问题。 
长河君 · 2026/9/1 2:22:27 · 阅读 1.2k 手撕Transformer：PyTorch实现与训练实战指南
手撕 Transformer 这件事很多人一开始会觉得特别难尤其是看到 QKV、多头注意力、位置编码这些名词堆在一起时很容易劝退。但等你真正把一个可运行的 PyTorch 代码从头写到尾再看到 loss 一点点降下去会发现 Transformer 的骨架其实非常清晰。本文就按照“概念 → 原理 → 代码 → 训练 → 排错”的顺序带你把 Transformer 从 Token 到训练完整过一遍所有代码都能直接跑重点部分我会拆开揉碎讲清楚。如果你正准备入门 NLP、做文本分类、做序列标注或者想搞懂 GPT 类的生成模型底层到底在干什么这篇文章都非常适合。你不需要有很深的深度学习基础只要会 Python装好了 PyTorch能跟着敲代码就足够了。1. 背景与核心概念为什么最后是 Transformer1.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 NetworkFFN。每个子层外面都套了残差连接和层归一化。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 python3.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 matplotlib2.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_embeddings10000, embedding_dim128) 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, 2i1) 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, dtypetorch.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 positions3.3 自注意力机制与 QKV自注意力是整个 Transformer 的灵魂。它的目的是让序列中的每个 Token 都能“关注”到其他 Token并根据关注程度聚合信息。先解释 QKV 这三个字母的含义QQuery查询向量表示当前 Token 想要从其他 Token 那里获取什么信息。KKey键向量表示当前 Token 能够提供什么信息。VValue值向量表示当前 Token 携带的实际内容。可以类比搜索引擎你输入搜索词Query系统匹配网页标题Key然后返回网页内容Value。计算过程分三步第一步通过三个线性变换把输入向量x映射成 Q、K、VQ 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, maskNone): 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, dropout0.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, maskNone): 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, dropout0.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, maskNone): # 第一个子层多头注意力 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 前馈网络 FFNFFN 是一个两层的全连接网络中间接一个 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_model128 时d_ff512。4. 完整实战手写一个可训练的 Transformer 分类模型这一节我们把前面拆开的模块组装起来做一个文本分类 Demo。数据集用一个小规模的英文情感分类数据包含正面和负面句子。模型结构是 Encoder-only 的 Transformer最后接一个全连接分类头。4.1 整体代码结构与流程整个项目只有一个 Python 文件建议命名为transformer_demo.py。流程如下构建一个小型数据集文本 Token 化并转换成 Token ID定义 Transformer 模型定义损失函数、优化器训练并输出 loss 和准确率验证模型效果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, dtypetorch.long)这里max_len取数据集中最长句子的长度。为了避免 padding 过多我们统一设置max_len8。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, dropout0.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, maskNone): # 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, epochs20, batch_size16, lr0.001): optimizer torch.optim.Adam(model.parameters(), lrlr) 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:ibatch_size] inputs torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels torch.tensor([label for _, label in batch], dtypetorch.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(fEpoch {epoch1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}) return losses4.6 跑起来vocab build_vocab(all_data) model TransformerEncoder( vocab_sizelen(vocab), d_model64, n_heads4, d_ff256, n_layers2, max_len8, num_classes2, dropout0.1, ) losses train_model(model, all_data, vocab, max_len8, epochs20)模型参数量非常小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, dtypetorch.long) # ---------- 3. 位置编码 ---------- def sinusoidal_position_encoding(seq_len, d_model): pe torch.zeros(seq_len, d_model) position torch.arange(0, seq_len, dtypetorch.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, dropout0.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, maskNone): 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, dropout0.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, maskNone): 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, dropout0.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, maskNone): 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, epochs20, batch_size16, lr0.001): optimizer torch.optim.Adam(model.parameters(), lrlr) 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:ibatch_size] inputs torch.stack([encode_sentence(sent, vocab, max_len) for sent, _ in batch]) labels torch.tensor([label for _, label in batch], dtypetorch.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(fEpoch {epoch1}/{epochs} - Loss: {avg_loss:.4f} - Acc: {acc:.4f}) # ---------- 8. 启动训练 ---------- vocab build_vocab(all_data) print(词表大小:, len(vocab)) model TransformerEncoder( vocab_sizelen(vocab), d_model64, n_heads4, d_ff256, n_layers2, max_len8, num_classes2, dropout0.1, ) print(模型参数总量:, sum(p.numel() for p in model.parameters())) train_model(model, all_data, vocab, max_len8, epochs20)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), diagonal1).bool() return ~mask第二生成是逐步进行的。训练时我们可以用 Teacher Forcing 一次性把完整目标序列输入模型但推理时模型只能先预测第一个 Token把它拼到输入后面再预测第二个 Token依次往复。如果未来你想做机器翻译或者写一个 GPT 风格的小模型可以在本文 Encoder 代码的基础上加入因果掩码把分类头换成词表大小的线性层然后使用交叉熵损失对下一个 Token 进行预测。本文的代码已经把自注意力、位置编码这些核心模块都实现了扩展成 Decoder 的难度并不会太大。6. 常见问题与排查思路前面代码虽然能跑通但你在自己写或者调整的时候大概率还会碰到各种问题。下面列出高频问题。问题现象常见原因解决思路loss 不下降学习率过大或过小尝试 1e-4 到 1e-2 之间的学习率观察前几个 step 的 loss 变化loss 为 NaN梯度爆炸或数据含 NaN检查输入是否含 NaN添加梯度裁剪clip_grad_norm_(model.parameters(), 1.0)准确率一直 50%模型结构 bug 或标签不平衡打印模型的 forward 输出形状确认最后一层输出和标签维度匹配训练很慢没有用 GPU 或 batch 太大确认torch.cuda.is_available()或者减小 batch size位置编码加到输入后结果不对位置编码没有转移到 GPU用.to(x.device)同步位置编码和输入张量注意力矩阵维度不匹配Q/K/V 拆头维度写错重点检查view和transpose后的形状建议每步都打印 shape安装 PyTorch 太慢网络原因或源问题用国内镜像或直接去官方站下载本地 wheel 安装针对 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 size4 是头数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_model512、n_heads8或者 d_model768、n_heads12。这个数字不是随便定的它保证了每个头有足够的信息容量同时拆分后计算效率较高。第二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_idx0): 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 并没有想象中那么神秘。如果本文对你有帮助可以收藏备用以后调参或者面试复习的时候翻出来看看。也欢迎在评论区交流你在跑代码过程中遇到的具体问题。