从零手写Transformer:PyTorch实现注意力机制与位置编码

发布时间:2026/9/1 15:56:42
从零手写Transformer:PyTorch实现注意力机制与位置编码 这篇我们来做一个非常硬核的实操不用 HuggingFace不调transformers库只用 PyTorch 的基础组件从零手写一个可以训练的 Transformer。从 Token 化、词嵌入、位置编码到多头自注意力、残差连接、层归一化、前馈网络再到完整的训练循环和批量推理全部手撕一遍。为什么要做这件事因为当你真正落到代码层面把每个张量的维度写对、把每个掩码加对、把训练循环跑通之后你对 “Transformer 为什么能 work” 的理解会远远超过只看论文和调包。这篇文章全程使用 PyTorch 基础 API目标是把架构的每一块拆开再用代码拼起来。先给结论这个实操可以在 CPU 上完整跑通不需要高配 GPU代码量集中在 6 个模块左右唯一需要的第三方库就是 PyTorch。如果你想快速过一遍核心流程可以直接跳到第 5 章开始看代码如果你的 PyTorch 环境还没装好先看第 4 章。1. 核心能力速览能力项说明项目类型PyTorch 手写 Transformer 教学项目核心组件Token 化、词嵌入、正余弦位置编码、多头自注意力、残差连接 LayerNorm、前馈网络 FFN运行环境CPU 可完整跑通GPU 可选依赖库Python PyTorch不依赖 HuggingFace训练数据自建小型文本分类数据集无需外部下载支持任务文本分类Encoder、可扩展为序列生成Decoder接口化部署提供 FastAPI 封装思路方便接入批量文本处理代码规模核心类约 7 个总代码量约 300 行需要说明的是这里不涉及真实业务数据集的测试显存、训练时长等指标需要按本机环境验证。但在 CPU 上用一个小数据集把整个流程跑通验证模型输出是否正常是完全没有问题的。2. Transformer 架构拆解从 Token 到输出Transformer 的核心思想是“注意力机制”。输入文本经过 Token 化变成整数序列再通过 Embedding 变成向量序列加入位置编码后送入多个编码器层。每个编码器层内部包含两个子模块多头自注意力子层和前馈网络子层每个子层都带残差连接和层归一化。整个流程可以拆成 6 步Token 化把文本切分成字符或子词映射成词表中的整数 ID。词嵌入把整数 ID 查表转成稠密向量。位置编码给每个位置的向量加入位置信息让模型知道先后顺序。多头自注意力计算每个 token 与其他 token 的相关性通过 QKV 机制完成特征交互。前馈网络 残差 层归一化对注意力输出做非线性变换同时通过残差和归一化稳定训练。输出层根据任务选择分类头或生成头。本文的代码实现采用 Encoder 架构用于文本分类第 9 章会额外给出 Decoder 中因果掩码的写法方便扩展为生成模型。3. 适用场景与使用边界这个项目最直接的适合人群是正在学 Transformer 原理、想通过代码加深理解的读者。准备面试算法岗或大模型岗需要手撕 Transformer 的求职者。想在业务中用小型自训练文本分类器但又不想引入重型依赖的开发者。希望掌握 PyTorch 张量维度变换、注意力掩码、训练循环的进阶学习者。能解决的问题本地训练一个轻量级文本分类模型数据量在几千条以内时效果可控。完全脱离开源库方便二次改造比如换成自定义注意力、加入可学习位置编码、融合外部特征。代码结构清晰每一层都能单独导出测试适合做教学和实验。不适合什么场景处理超大词表或超长文本需要的是优化过的推理框架而不是手写 Transformer。需要大规模预训练效果手写网络从头训练并不现实建议基于成熟的预训练模型微调。对推理性能要求极高应使用 TensorRT、ONNX Runtime 等方案。安全与合规边界涉及自定义数据训练时请确认数据来源合法不包含个人隐私信息。如果后续将模型部署到生产环境要对输入内容做过滤和审计防止生成或分类结果被滥用。本文代码仅用于学习和实验验证不构成任何商用承诺。4. 环境准备与前置条件4.1 安装 PyTorch在开始之前先确认 Python 环境。建议使用 Python 3.9 及以上版本。然后用 pip 安装 PyTorch。CPU 版本安装pip install torchGPU 版本安装前先通过nvidia-smi查看显卡驱动支持的 CUDA 版本再到 PyTorch 官网选择对应的命令。这里不指定具体版本号因为 PyTorch 迭代很快安装最新稳定版通常最稳妥。验证安装python -c import torch; print(torch.__version__); print(torch.cuda.is_available())如果返回True说明 GPU 可用返回False不代表失败CPU 也可以完成本项目的全部代码。4.2 目录结构建议把代码按模块拆分形成一个清晰的工程目录transformer_from_scratch/ ├── data.py # 数据预处理与 Token 化 ├── model.py # Transformer 模型定义 ├── train.py # 训练与评估脚本 ├── infer.py # 批量推理脚本 └── README.md目录拆分的目的是让每一部分职责清晰。实际项目可以按需合并。5. 数据预处理Token 化与词表构建5.1 什么是 TokenToken 是模型处理文本的最小单元。中文场景下可以直接按字符切分也可以使用分词工具切分成词。为了降低依赖本文采用字符级 Token 化实现简单效果在短文本分类上够用。分词函数def tokenize(text: str) - list[str]: # 中文按字符切分英文按空格切分 return list(text.replace( , ))以“电影真好看”为例切分结果[电, 影, 真, 好, 看]5.2 构建词表词表的作用是把 Token 映射成整数 ID。需要预留特殊符号pad用于批次内句子对齐。unk表示词表外 Token。cls用于分类任务标记句子开头。from collections import Counter def build_vocab(texts: list[str], min_freq: int 1) - dict[str, int]: counter Counter() for text in texts: counter.update(tokenize(text)) vocab {pad: 0, unk: 1, cls: 2} for token, freq in counter.items(): if freq min_freq and token not in vocab: vocab[token] len(vocab) return vocab # 使用示例 texts [电影真好看, 剧情很烂, 画面精美] vocab build_vocab(texts) print(vocab)输出示例{pad: 0, unk: 1, cls: 2, 电: 3, 影: 4, 真: 5, 好: 6, 看: 7, 剧: 8, 情: 9, 很: 10, 烂: 11, 画: 12, 面: 13, 精: 14, 美: 15}5.3 文本转 Tensor给定一段文本先 tokenize再映射成 ID最后做 padding 对齐到固定长度。import torch from torch.utils.data import Dataset, DataLoader class TextDataset(Dataset): def __init__(self, texts, labels, vocab, max_len32): self.texts texts self.labels labels self.vocab vocab self.max_len max_len def __len__(self): return len(self.texts) def __getitem__(self, idx): tokens tokenize(self.texts[idx]) ids [self.vocab[cls]] # 分类标记 for t in tokens[: self.max_len - 2]: ids.append(self.vocab.get(t, self.vocab[unk])) ids ids [self.vocab[pad]] * (self.max_len - len(ids)) return torch.tensor(ids, dtypetorch.long), torch.tensor(self.labels[idx])5.4 构造一个可直接运行的小数据集为了演示完整的训练流程使用 6 条手工构造的样本做正向/负向文本分类texts [ 这部电影太棒了, 垃圾剧情浪费时间, 画面精美演员演技在线, 不好看很差劲, 笑点密集值得一看, 节奏拖沓让人困, ] labels [1, 0, 1, 0, 1, 0] # 1 正向0 负向 dataset TextDataset(texts, labels, vocab, max_len16) dataloader DataLoader(dataset, batch_size2, shuffleTrue) for batch in dataloader: print(batch[0].shape, batch[1].shape) breaktorch.Size([2, 16]) torch.Size([2])到这里数据管线已经打通。6. 位置编码与词嵌入6.1 为什么需要位置编码自注意力机制本身没有序列顺序概念即使把句子的 Token 随机打乱注意力计算结果也不会改变。想让模型感知“谁先谁后”必须额外加入位置信息。最简单的方式是给每个位置一个可学习的向量但这样无法处理训练中没见过的超长序列。正余弦位置编码则通过固定公式生成位置向量具备一定外推能力。6.2 正余弦位置编码公式位置编码的核心公式如下PE(pos, 2i) sin(pos / 10000^(2i / d_model))PE(pos, 2i1) cos(pos / 10000^(2i / d_model))其中pos是位置索引i是维度索引d_model是嵌入维度。6.3 PyTorch 实现import math import torch import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int 512): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len).unsqueeze(1) # (max_len, 1) div_term torch.exp( torch.arange(0, d_model, 2) * (-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, max_len, d_model) self.register_buffer(pe, pe) def forward(self, x: torch.Tensor) - torch.Tensor: # x: (batch, seq_len, d_model) return x self.pe[:, : x.size(1), :]register_buffer的作用是让位置编码矩阵随模型一起迁移到 GPU但不会参与梯度更新。6.4 词嵌入实现词嵌入是一个简单的查表操作class TokenEmbedding(nn.Module): def __init__(self, vocab_size: int, d_model: int): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) def forward(self, x: torch.Tensor) - torch.Tensor: # x: (batch, seq_len) return self.embedding(x) # (batch, seq_len, d_model)在 Transformer 的输入阶段先嵌入再叠加位置编码class TransformerEmbedding(nn.Module): def __init__(self, vocab_size: int, d_model: int, max_len: int 512): super().__init__() self.token_embedding TokenEmbedding(vocab_size, d_model) self.position_encoding PositionalEncoding(d_model, max_len) def forward(self, x: torch.Tensor) - torch.Tensor: return self.position_encoding(self.token_embedding(x))为什么正余弦位置编码能work因为 sin/cos 函数满足线性关系模型可以通过注意力矩阵学习到相对位置信息而不只是绝对位置。这也是论文中选择它而不是直接学习位置向量的重要原因。7. 手写多头自注意力QKV 全流程多头自注意力是整个 Transformer 最核心的部分。本章从 QKV 的计算开始逐步完成缩放点积注意力、多头拆分和掩码。7.1 QKV 是什么QQuery、KKey、VValue三个向量由输入分别做线性变换得到。Query表示当前 token 想找什么信息。Key表示当前 token 能提供什么信息。Value表示当前 token 的实际内容。注意力分数就是 Query 和 Key 的点积归一化后再加权求和 Value。import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float 0.1): super().__init__() assert d_model % n_heads 0, d_model 必须能被 n_heads 整除 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, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None): batch_size query.size(0) # 线性变换后拆分多头 Q self.W_q(query).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) K self.W_k(key).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) V self.W_v(value).view(batch_size, -1, 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) attn torch.softmax(scores, dim-1) attn self.dropout(attn) context torch.matmul(attn, V) # 合并多头 context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.W_o(context)7.2 关键维度分析以 batch2、seq_len16、d_model128、n_heads4 为例输入 x 的形状(2, 16, 128)经过W_q后仍是(2, 16, 128)view(batch_size, -1, n_heads, d_k)后(2, 16, 4, 32)transpose(1, 2)后(2, 4, 16, 32)注意力分数(2, 4, 16, 16)这里最关键的一步是transpose操作它把“序列维度”和“头维度”互换让每个头独立计算注意力。7.3 为什么除以 sqrt(d_k)当维度较高时点积结果会变得很大导致 softmax 输出落入梯度饱和区。除以sqrt(d_k)将方差拉回 1 附近训练更稳定。7.4 mask 的作用mask 有两种常见用途padding mask遮盖pad位置避免模型关注无效 Token。causal mask遮盖未来信息用于 Decoder 生成。生成 padding mask 的方法def make_padding_mask(token_ids: torch.Tensor, pad_id: int 0) - torch.Tensor: # token_ids: (batch, seq_len) mask (token_ids ! pad_id).unsqueeze(1).unsqueeze(2) # (batch, 1, 1, seq_len) return mask使用方式是在注意力分数做masked_fill将非法位置替换为负无穷。8. 搭建完整 Transformer 编码器多头注意力只是 Transformer 的一个子层。完整的编码器层还包含前馈网络、残差连接和层归一化。8.1 前馈网络 FFNFFN 是一个两层的全连接网络中间用 ReLU 激活并加 Dropoutclass PositionwiseFeedForward(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float 0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) def forward(self, x: torch.Tensor) - torch.Tensor: return self.linear2(self.dropout(torch.relu(self.linear1(x))))8.2 层归一化与残差连接每个子层都使用“残差 LayerNorm”结构。残差连接让梯度可以直接从深层回流到浅层LayerNorm 对特征维度做归一化稳定训练。x x dropout(sublayer(norm(x))) # 标准 Post-Norm 写法在本文中使用的是 Pre-LN 变体还是 Post-LN为了训练稳定性代码采用 Post-LN 的经典写法。实际工程里 Pre-LN 更常见训练更稳读者可以自行实验。8.3 编码器层class EncoderLayer(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float 0.1): super().__init__() self.self_attn MultiHeadAttention(d_model, n_heads, dropout) self.feed_forward PositionwiseFeedForward(d_model, d_ff, dropout) self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: # 子层 1多头自注意力 x x self.dropout(self.self_attn(x, x, x, mask)) x self.norm1(x) # 子层 2前馈网络 x x self.dropout(self.feed_forward(x)) x self.norm2(x) return x8.4 完整编码器模型把多个编码器层堆叠起来加上 Embedding、位置编码、池化层和分类头就是一个可以训练的 Transformer 分类器class TransformerEncoder(nn.Module): def __init__(self, vocab_size: int, d_model: int 128, n_heads: int 4, n_layers: int 2, d_ff: int 512, num_classes: int 2, max_len: int 64, dropout: float 0.1): super().__init__() self.embedding TransformerEmbedding(vocab_size, d_model, max_len) self.encoder_layers nn.ModuleList([ EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers) ]) self.pooler nn.AdaptiveAvgPool1d(1) self.classifier nn.Linear(d_model, num_classes) self.dropout nn.Dropout(dropout) def forward(self, token_ids: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: x self.embedding(token_ids) # (batch, seq_len, d_model) for layer in self.encoder_layers: x layer(x, mask) # (batch, seq_len, d_model) x x.transpose(1, 2) # (batch, d_model, seq_len) x self.pooler(x).squeeze(-1) # (batch, d_model) return self.classifier(self.dropout(x))为何使用AdaptiveAvgPool1d因为每个 batch 内的序列可能长度不同虽然 padding 后长度一致但平均池化仍然是一种简单有效的全局特征聚合方式。换成CLS Token池化或最大池化也同样可行。9. 训练模型与效果验证9.1 定义损失函数、优化器与训练循环分类任务使用交叉熵损失优化器选择 AdamW并加一个简单的学习率调度器import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model TransformerEncoder( vocab_sizelen(vocab), d_model64, n_heads4, n_layers2, d_ff128, num_classes2, max_len16, dropout0.1 ).to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-3) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) def train_epoch(model, dataloader, optimizer, criterion, device): model.train() total_loss 0.0 for input_ids, labels in dataloader: input_ids, labels input_ids.to(device), labels.to(device) optimizer.zero_grad() logits model(input_ids) loss criterion(logits, labels) loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) def evaluate(model, dataloader, device): model.eval() correct 0 total 0 with torch.no_grad(): for input_ids, labels in dataloader: input_ids, labels input_ids.to(device), labels.to(device) logits model(input_ids) preds logits.argmax(dim-1) correct (preds labels).sum().item() total labels.size(0) return correct / total if total 0 else 0.09.2 执行训练训练 30 轮每轮打印 loss最后输出准确率epochs 30 for epoch in range(1, epochs 1): loss train_epoch(model, dataloader, optimizer, criterion, device) scheduler.step() if epoch % 5 0 or epoch 1: acc evaluate(model, dataloader, device) print(fEpoch {epoch:02d} | Loss: {loss:.4f} | Acc: {acc:.2f})在 6 条样本的小数据集上模型会很快过拟合训练集这是正常现象。目的是验证从 Token 到训练的全流程是否走通而不是评估真实泛化能力。判断训练成功的标准Loss 呈下降趋势不出现 NaN。训练集准确率达到 1.0。对新输入的推理结果符合预期。9.3 推理测试model.eval() test_texts [这部电影非常精彩, 烂片浪费时间] for text in test_texts: tokens tokenize(text) ids [vocab[cls]] [vocab.get(t, vocab[unk]) for t in tokens[:14]] ids torch.tensor([ids [vocab[pad]] * (16 - len(ids))]).to(device) with torch.no_grad(): logits model(ids) pred logits.argmax(dim-1).item() print(f文本: {text} - 预测类别: {pred})如果输出 1 表示正向输出 0 表示负向说明模型已经学到了基本的文本特征。9.4 扩展到生成任务Decoder 的因果注意力如果要把 Transformer 从分类扩展到文本生成核心区别在于 Decoder 使用因果掩码让当前位置只能看到历史信息。def make_causal_mask(seq_len: int, device: torch.device) - torch.Tensor: mask torch.triu(torch.ones(seq_len, seq_len), diagonal1).bool().to(device) return ~mask # False 表示相同位置及之前的位置可见 # 在注意力计算时使用 # scores scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0) 0, -1e9)加入因果掩码后自注意力变成“自回归式”模型在预测第 t 个 token 时只能看到 t-1 以及之前的 token。这就是 GPT 系列模型能逐 token 生成的核心机制。10. 批量推理与接口化部署模型训练好后可以封装成批量推理脚本再通过 HTTP 接口暴露服务。10.1 批量推理脚本def predict_batch(model, texts: list[str], vocab: dict[str, int], max_len: int 16) - list[int]: model.eval() batch_ids [] for text in texts: tokens tokenize(text) ids [vocab[cls]] [vocab.get(t, vocab[unk]) for t in tokens[: max_len - 1]] ids ids [vocab[pad]] * (max_len - len(ids)) batch_ids.append(ids) input_ids torch.tensor(batch_ids, dtypetorch.long).to(device) with torch.no_grad(): logits model(input_ids) preds logits.argmax(dim-1).cpu().tolist() return preds10.2 FastAPI 封装示例# app.py from fastapi import FastAPI from pydantic import BaseModel app FastAPI() class PredictRequest(BaseModel): texts: list[str] class PredictResponse(BaseModel): texts: list[str] labels: list[int] app.post(/predict, response_modelPredictResponse) def predict(req: PredictRequest): labels predict_batch(model, req.texts, vocab) return PredictResponse(textsreq.texts, labelslabels)启动服务uvicorn app:app --host 0.0.0.0 --port 8000调用接口curl -X POST http://127.0.0.1:8000/predict \ -H Content-Type: application/json \ -d {texts: [电影好看, 剧情很差]}{texts:[电影好看,剧情很差],labels:[1,0]}需要提醒的是生产环境部署时一定要加请求频率限制、输入长度校验、鉴权机制避免接口被滥用。11. 资源占用与性能观察11.1 CPU 与 GPU 差异在本文所使用的小模型配置下d_model64、n_layers2、n_heads4CPU 和 GPU 都能快速完成训练。真正拉大差距的是大词表、长序列、大 batch。观察资源占用的方法Linux/Mac 上用top或htop观察 CPU 和内存。Windows 用任务管理器看内存占用。如果有 GPU用nvidia-smi观察显存。nvidia-smi关注Memory-Usage一列可以看到 PyTorch 分配了多少显存。11.2 影响资源占用的因素序列长度注意力矩阵大小是seq_len * seq_len序列长度翻倍显存占用约翻 4 倍。batch_sizebatch 越大显存占用线性增长。d_model模型宽度越大参数和中间激活越多。n_heads多头数量主要影响注意力矩阵的拆分方式对参数量增加有限。总训练步数影响的是时间不是峰值显存。11.3 降低显存占用的常用方法使用梯度累积用小 batch 模拟大 batch。使用混合精度训练torch.cuda.amp。降低max_len限制序列长度。使用 Pre-LN 的 Transformer 变体训练更稳定对大模型更友好。11.4 进程残留问题训练时如果使用 GPU程序意外退出后显存可能未释放。在 PyTorch 中执行python -c import torch; print(torch.cuda.memory_summary())如果发现 GPU 显存被占用可以通过进程管理工具找到残留进程并终止。12. 常见问题与排查方法问题现象可能原因排查方式解决方案PyTorch 安装后无法导入Python 版本不匹配或安装命令错误python --version检查 pip 源使用官方源重装确认 Python 3.9CUDA 不可用显卡驱动版本过旧或安装了 CPU 版nvidia-smi查看驱动版本更新驱动重新安装对应 CUDA 版 PyTorch运行时维度不匹配多头注意力的 d_model 没有整除 n_heads打印张量 shape 核对调整 d_model 或 n_headsLoss 不下降学习率过大或过小词表映射错误打印梯度、检查 vocab调低学习率确认 Token ID 映射正确Loss 出现 NaN注意力分数未做 mask被 pad 位置污染检查 mask 是否生效在 softmax 前对非法位置填 -1e9推理结果全为同一类别数据不均衡或模型过拟合打印各类别数量增加数据量或用加权损失训练后模型无法保存路径不存在检查 save 路径先创建目录再保存FastAPI 启动失败端口被占用或未安装 uvicornnetstat -ano查看端口换端口或pip install uvicorn补充两个常见安装问题PyTorch 下载速度很慢时可以使用国内镜像源pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simpleWindows 下安装 GPU 版 PyTorch 时如果torch.cuda.is_available()返回 False大概率是 CUDA 版本与 PyTorch 要求不匹配。先在命令行执行nvidia-smi查看右上角 CUDA Version然后选择不高于该版本的 PyTorch CUDA 版本安装。13. 最佳实践与使用建议13.1 先小后大第一次跑训练时把d_model设为 32n_layers设为 1max_len设为 16batch_size 设为 2。确认代码跑通、Loss 正常下降后再逐步调大模型规模。这样可以从一开始就排除大部分代码问题。13.2 模块化测试每写完一个模块单独打印输出的 tensor shapeprint(f输入: {input_ids.shape}) print(f词嵌入: {embedded.shape}) print(f位置编码后: {encoded.shape}) print(f注意力输出: {attn_out.shape})维度不匹配是手写网络最常见的问题随时打印 shape 是最快的排错方法。13.3 保存最小可运行配置把模型参数、模型状态、优化器状态和词表写到同一个目录import os save_dir ./checkpoint os.makedirs(save_dir, exist_okTrue) torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), vocab: vocab, config: { vocab_size: len(vocab), d_model: 64, n_heads: 4, n_layers: 2, d_ff: 128, max_len: 16, num_classes: 2, } }, os.path.join(save_dir, model.pt))加载时先根据 config 重建模型再加载 state_dict避免参数不匹配。13.4 数据与代码分离输入文本、词表、模型输出分目录管理data/ # 原始数据 models/ # 训练好的模型 checkpoints/ # 训练中间状态 outputs/ # 推理结果 logs/ # 训练日志13.5 接口服务的安全性如果要把模型部署成 HTTP 服务至少做这几件事限制单次请求的文本条数和单条长度。增加 API Token 鉴权。对输入文本做敏感词过滤。记录请求日志方便审计和回滚。13.6 数据合规提醒训练数据如果来自公开文本或用户内容必须确认版权归属和数据使用协议。涉及个人信息的文本要做脱敏处理。涉及人脸、声音、肖像的内容更要确认授权后再使用。14. 总结与下一步这次我们从零开始手写了一个可以训练的 Transformer包括字符级 Token 化、词表构建、正余弦位置编码、词嵌入、多头自注意力、残差连接、层归一化、前馈网络、训练循环、批量推理和 FastAPI 接口封装。这个小项目最值得尝试的点是它完全脱离了transformers库的封装把注意力机制、QKV 计算、mask 处理这些核心概念全部暴露在你能直接看到的代码里。最先应该验证的功能是用第 5 章的数据集跑通 30 轮训练确认 Loss 下降且推理结果正确。最容易踩的坑是维度不匹配和 mask 未生效建议在每一步打印张量 shape。后续可以继续扩展的方向有三个。第一个是把 Decoder 补全加上因果掩码和 Cross-Attention做成一个标准的 Seq2Seq 模型用于文本生成或机器翻译。第二个是加入可学习位置编码、Pre-LN、SwiGLU 等现代改良对比不同配置的效果变化。第三个是把模型接入实际业务数据用几千条真实标注训练一个轻量分类器再用 FastAPI 部署成内部服务。如果感兴趣可以先保存这份代码框架后面按自己的数据改一版出来跑一下再回头对照论文细节理解会深很多。