Transformer 从零实现完整教学(零基础极致细化版)
本文的目标读者是完全零基础的学习者:不懂 Transformer、不太会深度学习、甚至没写过 PyTorch 都可以读懂。 全文按 数学准备 → 动机 → 原理手算推导 → NumPy 实现 → PyTorch 逐行实现 → 完整训练 → 官方框架 → 调试 的顺序展开, 每一步都有数值示例、形状标注、可运行代码,并与本仓库两个信号识别项目的真实代码相互印证:
- meta-transformer-amc-main/models/vit.py —— 教科书级多头注意力实现
- MCDformer-main/models/MCDformer.py —— timm 风格 Transformer Block 实现
姊妹文档:《Transformer特化方法与实践.md》(如何把标准 Transformer 改造成领域专用架构)。
目录
- 第 0 章 零基础准备:最少必要知识
- 第 1 章 动机:为什么要发明 Transformer
- 第 2 章 原理篇:一步一步推导注意力
- 第 3 章 代码篇:从零实现
- 第 4 章 框架篇:torch.nn 官方实现
- 第 5 章 调试篇:报错与陷阱
- 第 6 章 小结与下一步
第 0 章 零基础准备:必要知识
0.1 数学准备(简单数学 + 线性代数)
0.1.1 向量与点积
向量就是一行(或一列)数字。例如一个 4 维向量:
v = [1, 2, 3, 4]点积(内积):两个长度相同的向量,对应位置相乘再求和:
例:
点积的直觉:它衡量两个向量的方向相似度。方向越一致,点积越大;正交时点积为 0;方向相反时点积为负。整篇 Transformer 最核心的一步 Q·K 就是点积——“这个查询和那个键有多相似”。
0.1.2 矩阵乘法 = “行的加权组合”
矩阵 A 乘矩阵 B(记为 A @ B 或 AB)的规则:A 的列数必须等于 B 的行数。
只需要理解一点就可以了:设 A 是 m×k 矩阵、B 是 k×n 矩阵,则结果 C 是 m×n 矩阵,且 C 的第 i 行 = B 的所有行按 A 第 i 行的数字加权求和。
数值示例:
验证”行的加权组合”视角:C 第 1 行 = 1×[5,6] + 2×[7,8] = [5+14, 6+16] = [19, 22] ✓
这个视角为什么重要:注意力最后一步 A @ V 就是”用注意力权重矩阵 A 的每一行,去加权组合 V 的每一行”。V 的行 = 各个 token 的”值”,A 的行 = 各 token 对所有人的关注度。
token:模型处理的最小序列单元。NLP 里是一个词(或子词),本文信号场景里是”一个时间片/patch”。一句话 = 一串 token,注意力就是在这些 token 之间两两对话。
0.1.3 转置
转置(符号 ^T)把矩阵行列互换:B^T 的第 i 行第 j 列 = B 的第 j 行第 i 列。
在代码里写作 B.transpose(-2, -1)(交换最后两维)或 B.T(二维时)。
0.1.4 广播(broadcasting)
形状不同的张量相加时,PyTorch 会自动”复制”小张量以对齐形状。例如位置编码 pe 形状 (1, N, d),token 嵌入 x 形状 (B, N, d),x + pe 时 pe 沿 batch 维自动复制 B 份。记住:形状 (1, N, d) 的张量可以加到 (B, N, d) 上,这是位置编码的标准用法。
0.1.5 指数、对数与 softmax
(代码里 exp(x)):以 为底的指数函数,恒为正,且增长极快。
: 的反函数,。
softmax 把任意一列数字变成”概率分布”(每一项在 0~1 之间、总和为 1):
例:z = [2, 1, 0] → exp = [7.389, 2.718, 1] → 总和 11.107 → softmax = [0.665, 0.245, 0.090]。
作用:① 把打分变成权重;② 大数被指数放大,实现”赢者通吃”式的聚焦。注意力里 softmax 把 Q·K 的相似度打分变成”权重和为 1 的关注度”。
0.1.6 方差与均值(为什么除以 √d 会用到)
均值:。方差:,衡量数据”散开”的程度。方差越大,数值越可能偏离 0 很远。后文 2.4 节用它解释注意力里的缩放。
0.2 深度学习准备(三个核心概念)
- 模型 = 带参数的函数。神经网络就是一个函数
,是几百万个可调数字(参数)。训练的目标是找到让"预测"最接近"答案"的那组。 - 损失函数衡量预测有多差。例如分类用交叉熵损失(2.12 节):预测错了损失大,预测对了损失接近 0。
- 梯度下降更新参数:。 是损失对每个参数的”敏感度”(这个参数变大一点,损失会变多少),由反向传播(链式法则)自动算出; 是学习率(步长)。PyTorch 里这两步就是
loss.backward()和optimizer.step()。
两个实用概念:
- 过拟合:模型把训练数据背下来了,换新数据就崩。对策:正则化、dropout(训练时随机丢弃一部分神经元)、更多数据;
- dropout:训练时以概率 p 随机把一些数值清零,迫使模型不能依赖任何单一特征。测试时不丢弃。
0.3 PyTorch 准备
0.3.1 Tensor(张量)
PyTorch 的多维数组,与 NumPy 数组几乎一一对应:
import torch
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # 形状 (2, 2)print(x.shape) # torch.Size([2, 2])print(x[0, 1]) # 2.0 —— 索引print(x @ x) # 矩阵乘法print(x.transpose(0, 1))# 转置
a = torch.randn(2, 3) # 标准正态随机数,形状 (2, 3)b = torch.zeros(1, 3) # 全 0,形状 (1, 3)print((a + b).shape) # (2, 3) —— 广播形状约定:本文统一使用 (B, N, d) 表示 —— B = batch(一次处理的样本数)、N = 序列长度(token 数)、d = 每个 token 的特征维度。看任何张量先问三件事:哪一维是 batch?哪一维是序列?哪一维是特征?
0.3.2 自动求导
x = torch.tensor(3.0, requires_grad=True) # 声明"我要对它求导"y = x ** 2 # y = x²y.backward() # 计算 dy/dxprint(x.grad) # 6.0 —— 即 2x|x=3深度网络里所有参数的 .grad 就是这样被 loss.backward() 一次填好的。
0.3.3 nn.Module 与 nn.Linear(所有模块的模板)
import torch.nn as nn
class MyModel(nn.Module): # 所有模型都继承 nn.Module def __init__(self): super().__init__() self.fc = nn.Linear(4, 2) # 线性层: y = xW^T + b, 4 入 2 出 # W 形状 (2,4), b 形状 (2,), 自动注册为模型参数并参与梯度
def forward(self, x): # 定义前向计算 return self.fc(x)
model = MyModel()x = torch.randn(8, 4) # (B=8, 4)print(model(x).shape) # (8, 2)print(sum(p.numel() for p in model.parameters())) # 参数量: 2*4+2=10线性层是深度学习的”积木”:y = x W^T + b,即”对输入做一次加权求和”。注意力里的 Q/K/V 投影、FFN、分类头全都是线性层。nn.Parameter 则用来声明”裸参数”(如可学习位置编码)。
0.3.4 最小训练循环(10 行看懂整个训练流程)
import torch, torch.nn as nn
x = torch.randn(100, 4)y = (x.sum(dim=1) > 0).long() # 假任务: 和 >0 则类别 1
model = nn.Linear(4, 2)opt = torch.optim.Adam(model.parameters(), lr=0.01) # 优化器loss_fn = nn.CrossEntropyLoss() # 损失函数
for step in range(200): logits = model(x) # ① 前向: 预测 loss = loss_fn(logits, y) # ② 算损失 opt.zero_grad() # ③ 清空旧梯度(否则会累加) loss.backward() # ④ 反向传播: 填满 .grad opt.step() # ⑤ 按梯度更新参数
print("acc =", (model(x).argmax(1) == y).float().mean().item())# 输出接近 1.0,说明训练成功这个 ①→⑤ 循环就是所有深度学习训练的全部。后文第 3 章的完整 Transformer 训练只是把它放大。先弄清三个术语:logits 是模型输出的原始分数(未经 softmax 的每个类别打分,损失函数内部会自动归一化);epoch = 把全部数据过一遍网络;batch = 一次喂给模型的一小批数据;step = 每处理一个 batch 更新一次参数(上面代码每次 step 都用全部 100 条数据,即 batch = 整个数据集,200 个 step 就是把数据反复用了 200 遍)。注意 argmax(1):沿第 1 维(类别维)取最大值的下标,即预测类别。
0.4 环境安装与验证
pip install torch numpy验证:
import torchprint(torch.__version__) # 例如 2.4.0print("CUDA:", torch.cuda.is_available()) # True 表示有 GPU本文全部代码默认在 GPU 上运行:第 3 章开头用 torch.set_default_device('cuda') 让之后创建的所有张量/模块自动落在 GPU 上,正文代码无需手动写 .to("cuda")(训练演示在 GPU 上几十秒内完成)。
0.5 本文代码约定
- 每个代码块可独立复制运行(除了明确标注”片段”的);
- 关键张量旁标注形状注释
# (B, N, d); - 随机种子统一
torch.manual_seed(0)保证可复现; - 本文与两个项目的对照位置会给出可点击的相对路径。
第 1 章 动机:为什么要发明 Transformer
1.1 序列建模与词序问题
机器翻译任务:输入 “I love you”,输出 “我爱你”。两个难点:
- 序列长度不固定:输入 3 个词,输出 3 个字,但别的句子长短不一;
- 词序对齐复杂:英语的主谓宾和中文并不一一对应,有时词序颠倒、有时一对多(“爱” 对应 “love”)。
2017 年之前的主流方案是 RNN(循环神经网络)/LSTM(长短期记忆网络,RNN 的改进版,用”门”机制缓解遗忘):从左到右逐个词读入,把历史压缩进一个隐状态。但它有三个硬伤:
| 问题 | 后果 |
|---|---|
| 串行计算 | 第 t 步必须等第 t-1 步,GPU 无法并行,长序列训练极慢 |
| 长程依赖差 | 信息每传一步就”稀释”一次,第 100 个词很难记住第 1 个词(梯度消失) |
| 信息路径长 | 位置 i 的信息要经过 |i−j| 步才能影响位置 j |
1.2 注意力思想:让所有位置直接”开会”
Transformer 的答案简单粗暴:让任意两个位置直接连线。每个词向全句所有词提问、听取回答、更新自己。信息传播路径 O(1),全部计算可并行。
直觉:RNN 像”传话游戏”——一句话从左传到右,传到最后早就走样了。注意力像”圆桌会议”——每个人直接听所有人发言,自己决定采信多少。
1.3 Transformer 家族全景(你将要实现的属于哪一类)
| 家族 | 代表 | 结构 | 用途 | 本仓库对应 |
|---|---|---|---|---|
| Encoder-Decoder | 原版 Transformer | 完整编解码 | 机器翻译 | 本文第 3.9 节实现 |
| Encoder-only | BERT、ViT | 只有编码器 + 任务头 | 分类/理解 | meta-transformer 的 ViT、MCDformer、CTDNN |
| Decoder-only | GPT | 只有解码器(因果掩码) | 文本生成 | 无 |
两个项目都是 Encoder-only + 自定义任务头,所以本文主线是 Encoder(第 3.23.7 节),Decoder 作为完整实现的一部分(第 3.83.9 节)。
第 2 章 原理篇:一步一步推导注意力
本章用具体的数字手算一遍注意力的每一步。看不懂公式没关系,跟着数字算一遍就懂了。
2.1 检索类比:Query / Key / Value
想象你在图书馆找资料(Query:你的需求),每本书侧面贴着标签(Key:书的内容摘要)。你拿需求与每个标签比相似度,得到每本书的”重要度”,再按重要度把书的内容(Value)加权混合。
写成一行公式(这就是论文里的 Scaled Dot-Product Attention):
2.2 手算示例:2 个 token 的完整注意力
设序列只有 2 个 token,每个 2 维(d=2),输入:
X = [[1, 0], ← token 0 [0, 1]] ← token 1第 1 步:Q、K、V 从哪来? 为了先看清数学,让三个投影矩阵都是单位矩阵(即 Q=K=V=X,这个简化不会影响理解):
Q = K = V = X = [[1, 0], [0, 1]]第 2 步:算相似度打分 Q @ K^T(形状 (2,2),第 (i,j) 项 = token i 的 Q 与 token j 的 K 的点积):
第 3 步:缩放 除以 :
[[0.707, 0.000], [0.000, 0.707]]第 4 步:softmax(逐行!) 第 1 行 [0.707, 0]:
第 2 行对称,得到注意力权重矩阵:
第 5 步:加权求和 A @ V:
这个例子揭示了注意力的本质:每个 token 的输出是所有 token 的值的加权平均(权重矩阵每行和为 1,即”凸组合”)。注意力把信息从全序列”搬运”到每个位置——谁重要就多搬一点。训练的目的就是学习 W_Q, W_K, W_V,让模型知道”什么时候该关注谁”。
2.3 Q、K、V 的真正来源:三个投影矩阵
实际中 Q/K/V 不相等,而是用三个可学习的线性层把同一输入投影到三个”角色”:
- 输入 :N 个 token,每个 d 维;
- :可学习参数,把每个 token 投影成”查询/键/值”三种身份;
- Q 和 K 用来算相似度,V 是被搬运的内容。三权分立让模型既能灵活决定”关注谁”,又能灵活决定”搬运什么”。
“自注意力”(self-attention)的含义:Q、K、V 都来自同一序列自己(X 投影三次)——序列内部互相”开会”。
2.4 为什么除以 √d_k:方差分析(完整推导)
假设 q 和 k 的各分量独立、均值 0、方差 1。点积:
每一项 :均值 ;方差 。d 项相加后方差 。
也就是说 d 越大,点积的数值散布越广。而 softmax 是指数函数——输入稍大(比如 30 vs 3),,权重直接变成 one-hot,梯度接近 0(softmax 饱和),训练停滞。
除以 后方差回到 1,打分始终在”温和”区间,梯度健康。这就是 的来历——它不是魔法,是一次方差标准化。
补充直觉:维度越高,“向量恰好很像”或”恰好相反”的极端情况越容易发生,必须压制。
2.5 softmax 的数值稳定版本(代码里必须这么写)
直接算 遇到大数会溢出()。稳定做法:先减掉本行最大值再算:
减一个常数不改变 softmax 结果(分子分母同时除以 ),但保证最大输入是 0,exp 不会爆。PyTorch 的 F.softmax 内置了这个技巧,但你手写时要知道原理。
2.6 多头注意力:把 d 拆成 H 份并行”开会”
单头只有一个”关注模式”。多头把 d 维特征切成 H 组(每组 d_h = d/H 维),每组独立做一次完整注意力(有自己的 W_Q, W_K, W_V),最后拼起来再过一层线性投影:
形状全流程(这一张表背下来,代码永远不写错):
输入 X (B, N, d)投影 Q,K,V (B, N, d) 三个线性层各自拆头 reshape (B, N, H, d_h) → permute → (B, H, N, d_h)打分 Q@K^T (B, H, N, N) 每个头一张 N×N 注意力图softmax (B, H, N, N) 沿最后维(对 key)归一化加权 attn@V (B, H, N, d_h)拼头 transpose → reshape → (B, N, d)输出投影 W_O (B, N, d)B=batch,N=token 数,H=头数,d_h = d/H。头是”平行小注意力”,batch 是”同时处理的多个样本”——两者都是独立计算的,别混淆。
2.7 位置编码:注意力天生”看不见顺序”
观察 2.2 节:如果把 token 顺序打乱,注意力权重矩阵只是跟着行列置换,信息混合的方式完全不变——自注意力对顺序不敏感(数学上叫置换等变性)。但”我爱你”和”你爱我”显然不同,所以必须给每个位置注入位置信息。
2.7.1 正弦位置编码(原论文)
对位置 pos 的第 2i 维(偶数)和第 2i+1 维(奇数):
手算一个例子(d=4,位置 0 和 1):
pos=0: [sin(0), cos(0), sin(0), cos(0)] = [0, 1, 0, 1]pos=1: [sin(1), cos(1), sin(1/100), cos(1/100)] ≈ [0.841, 0.540, 0.010, 1.000]两个关键性质:
- 不同维度 = 不同波长: 当 i 从 0 到 d/2 变化时,波长从 (约 6.28)增长到 。低维(波长短)编码精细偏移,高维(波长长)编码大致位置——像二进制计数一样分层表示位置;
- 相对位置可线性表示: 可以用 乘一个旋转矩阵得到(旋转角只与 k 有关)。这意味着模型有可能”学会”直接比较相对距离。
实现要点:位置编码加到嵌入上(x = x + pe),不是拼接——保持 d 维不变。
2.7.2 可学习位置编码(ViT 与两个项目都用)
直接把位置编码声明为可训练参数,让数据自己学:
self.pos_embed = nn.Parameter(torch.zeros(1, N, d)) # 随训练更新优点:灵活、任务相关;缺点:长度写死(变长输入的处理见特化文档第 4 章)。
2.8 前馈网络 FFN:token 内部的”特征加工”
形状:d → 4d → d(中间隐层通常放大 4 倍)。两个关键理解:
- 逐 token 独立:同一个 MLP 对每个 token 分别作用(等价于核大小 1 的卷积),token 之间在这里不交互;
- 分工:注意力负责 token 之间的混合(token-mixing),FFN 负责 token 内部各特征维度的混合(channel-mixing)。交替堆叠 = 完整表示学习。GELU 是 ReLU 的平滑版,现代实现默认用 GELU。
2.9 残差连接与 LayerNorm(手算示例)
2.9.1 残差连接:给梯度一条”高速公路”
深网络的梯度要穿过几十层,逐层相乘会指数衰减(消失)。残差把”原始信号”直通给下一层,梯度也多了一条加法直通路径(对 x 的梯度至少为 1)。这就是为什么 Transformer 能堆几十上百层。
2.9.2 LayerNorm:把每行的数值”标准化”
对每个 token 自己(沿特征维 d)算均值方差并标准化:
手算示例:token 特征 x = [1, 2, 3, 4](d=4,注意是对这一行自己归一化):
mean = 2.5var = ((1-2.5)²+(2-2.5)²+(3-2.5)²+(4-2.5)²)/4 = 1.25标准化: [-1.342, -0.447, 0.447, 1.342] (γ=1, β=0 时输出即此)作用:让每层输入的数值尺度稳定,训练更快更稳。γ、β 是可学习参数(每维一个),允许模型学出”更合适的尺度”。与 BatchNorm 的区别:BatchNorm 沿 batch 维归一化(依赖一批数据的统计),LayerNorm 沿特征维(每样本独立)——对序列任务更合适,且与 batch size 无关。
2.9.3 两种组合顺序:Post-LN 与 Pre-LN
Pre-LN 让梯度路径更干净、对学习率不敏感,是 ViT 及两个项目的选择。本文实现用 Pre-LN。
2.10 掩码三兄弟
注意力允许任意屏蔽某些位置(屏蔽 = 打分置 -inf,softmax 后权重为 0):
| 掩码 | 屏蔽谁 | 为什么 | 形状 |
|---|---|---|---|
| Padding 掩码 | <pad> 填充位置 | 填充符没有信息,不该被关注 | (B, N) |
| 因果掩码(下三角) | 未来位置 j > i | 生成时不能”偷看未来” | (N, N) |
| 无掩码 | 谁也不屏蔽 | Encoder 双向自由开会 | — |
因果掩码矩阵(N=4):
[[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]] 位置 i 只能看到 j ≤ i2.11 完整 Encoder-Decoder 架构(原版全景)
输入序列 → 词嵌入 + 位置编码 │ ▼┌───────────────────────────────────┐│ Encoder ×N 层 ││ x = x + MultiHeadAttn(LN(x)) │ 双向:看整句│ x = x + FFN(LN(x)) │└───────────────────────────────────┘ │ (编码器的输出 = 对源句的"理解") ▼┌───────────────────────────────────┐│ Decoder ×N 层 ││ x = x + MaskedAttn(LN(x)) │ ① 掩码自注意力:只看已生成的部分│ x = x + CrossAttn(LN(x)) │ ② 交叉注意力:Q 来自解码器, K/V 来自编码器│ x = x + FFN(LN(x)) │└───────────────────────────────────┘ │ ▼Linear + Softmax → 词表概率分布(下一个词)Encoder-only 模型(本文主线、两个项目所用)就是:嵌入 + 位置编码 → N 层 Encoder Block → 读出(cls/池化)→ 分类头。
2.12 训练三件套:损失、优化器、学习率调度
2.12.1 交叉熵损失(分类标配)
模型把正确类别的概率 p 预测得越接近 1,损失越小。PyTorch 里 nn.CrossEntropyLoss() 自动包含 log-softmax(所以模型输出原始 logits 即可,不要再手动 softmax;logits 即模型最后一层输出的原始分数,未归一化、可正可负,过 softmax 之后才变成概率)。
2.12.2 标签平滑(Label Smoothing)
把 one-hot 标签 换成 (K = 类别数, 常取 0.1)。例如 3 分类、:[1,0,0] → [0.933, 0.033, 0.033]。作用:禁止模型对预测过度自信(否则正确类 logit 会被推向无穷大),提升泛化。代码见 3.10 节。
2.12.3 Adam 优化器
带”惯性 + 自适应步长”的梯度下降,是 Transformer 的默认选择:torch.optim.Adam(params, lr=1e-3)。
2.12.4 学习率预热(Warmup)
原论文的 Noam 调度:
翻译成白话:前 warmup 步线性爬升,之后按 √步数 衰减。为什么需要预热:训练初期参数是随机的,梯度很大,一上来就用大学习率会把模型”踢飞”(第 5.2 节的 NaN 大多源于此)。完整代码见 3.10 节。
第 3 章 代码篇:从零实现
3.0 实现路线图
| 步骤 | 模块 | 对应原理 | 依赖 |
|---|---|---|---|
| 3.1 | NumPy 版注意力 | 2.2/2.6 | 仅 numpy |
| 3.2 | InputEmbedding | 0.3.3 | torch |
| 3.3 | PositionalEncoding / LearnablePositionalEncoding | 2.7 | torch |
| 3.4 | MultiHeadAttention | 2.2~2.6、2.10 | torch |
| 3.5 | PositionwiseFeedForward | 2.8 | torch |
| 3.6 | EncoderBlock | 2.9、2.11 | 3.4+3.5 |
| 3.7 | TransformerEncoder | 2.11 | 3.6 |
| 3.8 | DecoderBlock + 因果掩码 | 2.10、2.11 | 3.4+3.5 |
| 3.9 | 完整 Encoder-Decoder | 2.11 | 3.7+3.8 |
| 3.10 | 训练脚本(warmup + label smoothing) | 2.12 | 3.9 |
| 3.11~3.13 | 三个演示实验 | — | 3.7 |
3.1 第 0 步:用 NumPy 实现注意力(先彻底理解数学)
在碰 PyTorch 之前,用 NumPy 手写一遍 2.2 节的五个步骤:
import numpy as np
def softmax(z, axis=-1): """数值稳定版 softmax(2.5 节):先减最大值再指数。""" z = z - z.max(axis=axis, keepdims=True) # 防止 exp 溢出 e = np.exp(z) return e / e.sum(axis=axis, keepdims=True)
def attention_numpy(X, Wq, Wk, Wv, d_k): """X: (N, d) 一个序列。返回注意力输出与权重矩阵。""" Q = X @ Wq # (N, d_k) —— 投影出查询 K = X @ Wk # (N, d_k) —— 投影出键 V = X @ Wv # (N, d_k) —— 投影出值 scores = Q @ K.T / np.sqrt(d_k) # (N, N) —— 相似度打分并缩放 A = softmax(scores, axis=1) # (N, N) —— 逐行 softmax out = A @ V # (N, d_k) —— 加权组合 return out, A
# ---- 复现 2.2 节的例子 ----X = np.array([[1.0, 0.0], [0.0, 1.0]]) # 2 个 token, 每 token 2 维I = np.eye(2) # 单位矩阵: Q=K=V=Xout, A = attention_numpy(X, I, I, I, d_k=2)
print("注意力权重 A =\n", np.round(A, 3))print("输出 =\n", np.round(out, 3))运行输出(对照 2.2 节的手算结果):
注意力权重 A = [[0.67 0.33 ] [0.33 0.67 ]]输出 = [[0.67 0.33 ] [0.33 0.67 ]]验证了手算:每个 token 的输出 = 全体 token 值的加权平均。现在把这 20 行翻译成 PyTorch 的可训练模块。
3.2 模块 1:输入嵌入 InputEmbedding
把”离散 token”(词的编号、信号的采样值)变成 d 维向量:
import mathimport torchimport torch.nn as nnimport torch.nn.functional as F
torch.manual_seed(0) # 固定随机种子,保证结果可复现
# ---- 设备选择:默认 GPU ----# set_default_device 让之后创建的所有张量/模块自动落在 GPU 上,正文代码无需再写 .to()/.cuda()torch.set_default_device('cuda')print("默认设备: cuda ——", torch.cuda.get_device_name(0))
class InputEmbedding(nn.Module): """两种嵌入方式: vocab 不为 None: 查表嵌入(NLP 用)——token 是词的编号; vocab 为 None : 线性嵌入(信号/数值用)——token 是连续数值。 """
def __init__(self, d_model: int, vocab: int = None, in_features: int = 1): super().__init__() self.d_model = d_model if vocab is not None: self.embed = nn.Embedding(vocab, d_model) # 查表: (B,N) → (B,N,d) else: self.embed = nn.Linear(in_features, d_model) # 线性: (B,N,1) → (B,N,d)
def forward(self, x: torch.Tensor) -> torch.Tensor: return self.embed(x) * math.sqrt(self.d_model)
# ---- 使用示例 ----emb = InputEmbedding(d_model=16, vocab=1000) # NLP 风格ids = torch.randint(0, 1000, (2, 5)) # (B=2, N=5) 词编号print(emb(ids).shape) # (2, 5, 16)
emb2 = InputEmbedding(d_model=16, in_features=1) # 信号风格sig = torch.randn(2, 5, 1) # (B=2, N=5, 1) 标量序列print(emb2(sig).shape) # (2, 5, 16)逐行说明:
nn.Embedding(vocab, d):内部是一张(vocab, d)的查找表,输入词编号输出对应向量;nn.Linear(in_features, d):,把 1 维标量线性投影到 d 维(等价于”每个数值 × 一个 d 维向量”);- 乘以
√d_model:原论文的小技巧,防止嵌入值在加法中相对位置编码”太小”; - 本文演示用信号任务,所以主线用线性嵌入。
3.3 模块 2:位置编码(正弦 + 可学习,各配验证)
class PositionalEncoding(nn.Module): """正弦位置编码(2.7.1 节)。register_buffer 的内容不参与训练但会随模型保存。"""
def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): super().__init__() pe = torch.zeros(max_len, d_model) # (max_len, d) 先全 0 position = torch.arange(0, max_len).unsqueeze(1).float() # (max_len, 1): 0,1,2,... # 分母 10000^(2i/d):用 exp(log) 形式写,数值更稳(对 i=0,2,4,... 生成 d/2 个) div_term = torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) # (d/2,) pe[:, 0::2] = torch.sin(position * div_term) # 偶数维填 sin pe[:, 1::2] = torch.cos(position * div_term) # 奇数维填 cos self.register_buffer("pe", pe.unsqueeze(0)) # (1, max_len, d) 预备广播 self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N, d)。pe[:, :N] 取前 N 个位置,靠广播加到每个样本上 x = x + self.pe[:, : x.size(1), :] return self.dropout(x)
class LearnablePositionalEncoding(nn.Module): """可学习位置编码(2.7.2 节,ViT 与两个项目同款)。"""
def __init__(self, num_tokens: int, d_model: int, dropout: float = 0.1): super().__init__() self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, d_model)) self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, num_tokens, d) return self.dropout(x + self.pos_embed)
# ---- 验证:位置编码形状与数值 ----pe = PositionalEncoding(d_model=4, max_len=100)x = torch.zeros(2, 3, 4) # (B=2, N=3, d=4)out = pe(x)print(out.shape) # (2, 3, 4) —— 形状不变,只加信息print("位置 0 编码:", pe.pe[0, 0].tolist()) # [0, 1, 0, 1] —— 对照 2.7.1 手算print("位置 1 编码:", [round(v, 3) for v in pe.pe[0, 1].tolist()]) # ≈ [0.841, 0.540, 0.010, 1.0]逐行说明:
torch.arange(0, d_model, 2):[0, 2, 4, ...](步长 2 取偶数下标);pe[:, 0::2]:切片写法,0::2= 从 0 开始每隔 2 个取一列(偶数列);position * div_term:(max_len,1) × (d/2,)广播成(max_len, d/2)——每个位置 × 每个波长;register_buffervsnn.Parameter:buffer 不参与梯度(正弦编码是固定公式,没有可学参数),但会随state_dict保存/加载;Parameter 参与梯度(可学习位置编码要靠训练更新);- 验证输出与 2.7.1 节手算一致 ✓。
3.4 模块 3:多头注意力(全文最核心的代码,逐行讲解)
class MultiHeadAttention(nn.Module): """多头自注意力。实现方式:Q/K/V 合并为一个大线性层(原论文同款)。
形状流程(对照 2.6 节表格): x (B,N,d) → qkv (B,N,3d) → reshape (B,N,3,H,d_h) → permute (3,B,H,N,d_h) → scores (B,H,N,N) → softmax → attn@V → (B,H,N,d_h) → transpose+reshape (B,N,d) → proj (B,N,d) """
def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0, f"d_model({d_model}) 必须能被 n_heads({n_heads}) 整除" self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads # 每个头的维度 d_h self.qkv = nn.Linear(d_model, 3 * d_model) # 一次投影出 Q,K,V(省 3 次调用) self.proj = nn.Linear(d_model, d_model) # 输出投影 W_O self.attn_drop = nn.Dropout(dropout) self.proj_drop = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): B, N, d = x.shape
# ---- 第 1 步:投影 + 拆出 Q/K/V ---- # Linear 输出 (B,N,3d);reshape 成 (B,N,3,H,d_h); # permute(2,0,3,1,4) 把"第 2 维(3)"提到最前 → (3,B,H,N,d_h),方便拆包 qkv = ( self.qkv(x) .reshape(B, N, 3, self.n_heads, self.d_head) .permute(2, 0, 3, 1, 4) ) q, k, v = qkv[0], qkv[1], qkv[2] # 各 (B,H,N,d_h)
# ---- 第 2 步:打分 + 缩放 ---- # q @ k^T: (B,H,N,d_h)×(B,H,N,d_h)^T → (B,H,N,N) # 语义:每个头里,每个 query 与每个 key 的相似度 attn = (q @ k.transpose(-2, -1)) * (self.d_head ** -0.5)
# ---- 第 3 步:掩码(可选)---- # mask 为 0 的位置置 -inf;softmax(exp(-inf)=0) 后权重为 0 # 支持两种形状:(B,N) padding 掩码 → (B,1,1,N);(N,N) 因果掩码 → (1,1,N,N) if mask is not None: if mask.dim() == 2 and mask.shape == attn.shape[-2:]: mask = mask.unsqueeze(0).unsqueeze(0) # 方阵掩码:广播到 batch 与头 else: mask = mask.unsqueeze(1).unsqueeze(2) # 逐样本掩码:广播到头与 query attn = attn.masked_fill(mask == 0, float("-inf"))
# ---- 第 4 步:softmax 归一化(沿最后一维 = 对 key 求和为 1)---- attn = attn.softmax(dim=-1) attn = self.attn_drop(attn)
# ---- 第 5 步:加权求和 + 拼头 + 输出投影 ---- out = attn @ v # (B,H,N,d_h) out = out.transpose(1, 2).reshape(B, N, d) # (B,N,d):把头拼回特征维 out = self.proj_drop(self.proj(out))
if return_attn: return out, attn # 附带注意力图(可视化用) return out逐行精讲:
self.qkv = nn.Linear(d_model, 3*d_model):把三个投影合并成一个矩阵(原论文写法)。输出前 d 维是 Q、中间 d 维是 K、后 d 维是 V。也可以像 MCDformer.py 那样写成三个独立nn.Linear(d_model, d_model)(timm 风格),数学完全等价;reshape(B, N, 3, H, d_h):把 3d 维拆成”3 个角色 × H 个头 × d_h”;permute(2, 0, 3, 1, 4):重排维度顺序。(B,N,3,H,d_h) → (3,B,H,N,d_h),然后qkv[0]取到全部 batch 全部头的 Q;q @ k.transpose(-2, -1):transpose(-2,-1)交换最后两维 = 转置(0.1.3 节)。矩阵乘法自动对 batch 和头数两维并行(batch 是样本间独立,头是子空间间独立);(self.d_head ** -0.5):即1/√d_h(2.4 节);masked_fill(mask == 0, -inf):mask为 0 的位置打分变负无穷。-inf的 softmax = ✓。注意masked_fill会广播掩码:(B,N)的 padding 掩码要先unsqueeze(1).unsqueeze(2)成(B,1,1,N),(N,N)的因果掩码要unsqueeze(0).unsqueeze(0)成(1,1,N,N),才能对齐(B,H,N,N)的注意力矩阵;attn.softmax(dim=-1):沿最后一维(key 方向)归一化——每行权重和为 1;transpose(1,2).reshape(B,N,d):(B,H,N,d_h) → (B,N,H,d_h) → (B,N,d),头的输出按序拼接回 d 维;self.proj(W_O):让不同头的信息互相混合——没有它,各头输出只是”拼在一起”而非”融合”。
单元验证(形状 + 可训练性 + 掩码正确性):
mha = MultiHeadAttention(d_model=32, n_heads=4)x = torch.randn(2, 10, 32) # (B=2, N=10, d=32)out, attn = mha(x, return_attn=True)print("输出形状:", out.shape) # (2, 10, 32) —— 形状不变print("注意力图形状:", attn.shape) # (2, 4, 10, 10) —— 每头一张 10×10 图print("每行权重和:", attn[0, 0, 0].sum().item()) # ≈1.0 —— softmax 性质
# 可训练性检查:参数能收到梯度loss = out.sum()loss.backward()print("qkv 权重梯度非空:", mha.qkv.weight.grad is not None) # True
# 掩码检查:屏蔽第 0 个 key 后,第 0 列权重应为 0mask = torch.ones(2, 10); mask[:, 0] = 0 # (B,N) 形状的 padding 掩码out_m, attn_m = mha(x, mask=mask, return_attn=True)print("被屏蔽列权重:", attn_m[0, 0, :, 0].sum().item()) # 0.0 ✓3.5 模块 4:前馈网络 FFN
class PositionwiseFeedForward(nn.Module): """FFN(2.8 节): d → 4d → d。逐 token 独立,形状永不改变。"""
def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1): super().__init__() d_ff = d_ff or 4 * d_model self.fc1 = nn.Linear(d_model, d_ff) # 放大 4 倍给"加工空间" self.fc2 = nn.Linear(d_ff, d_model) # 缩回原维度 self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N, d) —— Linear 自动作用在最后一维,对每个 token 都一样 return self.fc2(self.dropout(F.gelu(self.fc1(x))))
ffn = PositionwiseFeedForward(32)print(ffn(torch.randn(2, 10, 32)).shape) # (2, 10, 32)3.6 模块 5:EncoderBlock(Pre-LN,现代主流)
class EncoderBlock(nn.Module): """Pre-LN 编码块(2.9.3 节): x = x + Attn(LN(x)) ← 先归一化再进子层,最后残差相加 x = x + FFN(LN(x)) """
def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1): super().__init__() self.norm1 = nn.LayerNorm(d_model) # 注意力前的 LN self.attn = MultiHeadAttention(d_model, n_heads, dropout) self.norm2 = nn.LayerNorm(d_model) # FFN 前的 LN self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): if return_attn: a_out, attn = self.attn(self.norm1(x), mask, return_attn=True) x = x + a_out # 残差:梯度直通 else: x = x + self.attn(self.norm1(x), mask) x = x + self.ffn(self.norm2(x)) # FFN 支路 + 残差 if return_attn: return x, attn return x
block = EncoderBlock(32, 4)x = torch.randn(2, 10, 32)out, attn = block(x, return_attn=True)print("块输出形状:", out.shape) # (2, 10, 32)print("块内注意力图:", attn.shape) # (2, 4, 10, 10)逐行说明:
nn.LayerNorm(d_model):对每个 token 的 d 维特征归一化(2.9.2 节),与 batch 无关;- 残差结构的意义再强调一次:
x + ...让梯度至少有一条”×1”的直通路。若训练时 loss 不降,第一件事就是检查残差是否写对; return_attn参数一路透传,供 3.12 节可视化——研究代码里很常见的写法(MCDformer 的Attention.forward里同样有register_hook参数用于捕获注意力图)。
3.7 模块 6:TransformerEncoder(N 层堆叠)
class TransformerEncoder(nn.Module): def __init__(self, n_layers: int, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1, block_cls: nn.Module = EncoderBlock): super().__init__() # ModuleList 才能被优化器"看见"全部参数;普通 list 不会注册参数! # block_cls 参数让消融实验可以注入无残差等变体块 self.layers = nn.ModuleList( [block_cls(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)] ) self.norm = nn.LayerNorm(d_model) # 末尾 LN(ViT 惯例)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): all_attns = [] for layer in self.layers: if return_attn: x, attn = layer(x, mask, return_attn=True) all_attns.append(attn) else: x = layer(x, mask) if return_attn: return self.norm(x), all_attns return self.norm(x)
encoder = TransformerEncoder(n_layers=2, d_model=32, n_heads=4)x = torch.randn(2, 10, 32)out, attns = encoder(x, return_attn=True)print("编码器输出:", out.shape) # (2, 10, 32)print("层数 × 每层注意力图:", len(attns), attns[0].shape) # 2 张 (2,4,10,10)逐行说明:
nn.ModuleList而非普通list:只有nn.Module容器(ModuleList/Sequential)里的模块才会被注册进模型参数表。用普通 list 会导致参数不更新——经典新手坑;- 末尾的
self.norm:ViT 的惯例(在读出前最后归一化一次)。BERT 也如此。对 2 层的玩具模型可省略,但保留无害。
3.8 模块 7:DecoderBlock 与因果掩码(生成任务的核心)
class DecoderBlock(nn.Module): """解码块(2.11 节)。与编码块的区别: ① 自注意力加因果掩码(只看过去); ② 多一个交叉注意力(Q 来自解码器,K/V 来自编码器输出)。"""
def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1): super().__init__() self.norm1 = nn.LayerNorm(d_model) self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) # ① 掩码自注意力 self.norm2 = nn.LayerNorm(d_model) self.cross_attn = CrossAttention(d_model, n_heads, dropout) # ② 交叉注意力 self.norm3 = nn.LayerNorm(d_model) self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
def forward(self, x: torch.Tensor, enc_out: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor: x = x + self.self_attn(self.norm1(x), causal_mask) # 只许看 j ≤ i x = x + self.cross_attn(self.norm2(x), enc_out) # 向编码器"提问" x = x + self.ffn(self.norm3(x)) return x
class CrossAttention(nn.Module): """交叉注意力:Q 来自 x(解码端),K/V 来自 enc(编码端)。 与自注意力的唯一区别:K/V 的投影层接收的是另一个张量。"""
def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0 self.d_model, self.n_heads = d_model, n_heads self.d_head = d_model // n_heads self.wq = nn.Linear(d_model, d_model) # Q 投影(用于解码端) self.wk = nn.Linear(d_model, d_model) # K 投影(用于编码端) self.wv = nn.Linear(d_model, d_model) # V 投影(用于编码端) self.proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, enc: torch.Tensor) -> torch.Tensor: B, N, d = x.shape E = enc.size(1) # 编码端 token 数(可与 N 不同) q = self.wq(x).reshape(B, N, self.n_heads, self.d_head).permute(0, 2, 1, 3) k = self.wk(enc).reshape(B, E, self.n_heads, self.d_head).permute(0, 2, 1, 3) v = self.wv(enc).reshape(B, E, self.n_heads, self.d_head).permute(0, 2, 1, 3) attn = (q @ k.transpose(-2, -1)) * (self.d_head ** -0.5) # (B,H,N,E) 注意是 N×E! attn = attn.softmax(dim=-1) out = (attn @ v).transpose(1, 2).reshape(B, N, d) return self.proj(self.dropout(out))
def make_causal_mask(n: int, device=None) -> torch.Tensor: """下三角矩阵(2.10 节):(i,j) 处 j ≤ i 才为 1。""" return torch.tril(torch.ones(n, n, device=device))
# ---- 验证 ----cm = make_causal_mask(4)print(cm) # 对照 2.10 节的 4×4 矩阵dec_block = DecoderBlock(32, 4)x = torch.randn(2, 5, 32) # 解码端 5 个 tokenenc = torch.randn(2, 10, 32) # 编码端 10 个 tokenout = dec_block(x, enc, make_causal_mask(5))print("解码块输出:", out.shape) # (2, 5, 32) —— 解码端 token 数不变逐行说明:
- 交叉注意力的注意力矩阵是
N×E(矩形):解码端每个 token 对编码端每个 token 打分。这是它与自注意力(N×N 方阵)在形状上的本质区别; make_causal_mask用torch.tril(下三角)一行生成。mask==0处被打分置-inf(3.4 节),实现”禁止偷看未来”;- 翻译任务里解码器输入是”已经生成的词”,编码器输出是”源句的理解”,交叉注意力就是两者之间的桥。
3.9 模块 8:完整 Encoder-Decoder 模型
class Transformer(nn.Module): """原版 Encoder-Decoder 结构(2.11 节全景图)。"""
def __init__(self, vocab_size: int, d_model: int = 32, n_heads: int = 4, n_enc: int = 2, n_dec: int = 2, d_ff: int = None, dropout: float = 0.1, max_len: int = 100): super().__init__() self.d_model = d_model self.src_embed = nn.Embedding(vocab_size, d_model) # 源词嵌入 self.tgt_embed = nn.Embedding(vocab_size, d_model) # 目标词嵌入 self.pos_enc = PositionalEncoding(d_model, max_len, dropout) # 共用一套位置编码 self.encoder = TransformerEncoder(n_enc, d_model, n_heads, d_ff, dropout) self.decoder = nn.ModuleList( [DecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_dec)] ) self.head = nn.Linear(d_model, vocab_size) # 词表概率
def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor: tgt_len = tgt.size(1) enc_out = self.encoder(self.pos_enc(self.src_embed(src))) # (B,S,d) x = self.pos_enc(self.tgt_embed(tgt)) # (B,T,d) causal = make_causal_mask(tgt_len, tgt.device) # (T,T) for blk in self.decoder: x = blk(x, enc_out, causal) return self.head(x) # (B,T,vocab)
# ---- 形状冒烟测试 ----model = Transformer(vocab_size=100)src = torch.randint(0, 100, (2, 8)) # 源句 8 词tgt = torch.randint(0, 100, (2, 5)) # 目标句 5 词logits = model(src, tgt)print("输出 logits 形状:", logits.shape) # (2, 5, 100) —— 每个位置预测下一个词逐行说明:
- 编码器与解码器共享同一套位置编码(都从位置 0 开始编码);
- 解码器输入是右移一位的目标句(训练时用”teacher forcing”:给模型看正确答案的前缀,让它预测下一个词);
- 输出
(B, T, vocab):每个已生成位置预测下一个词的分布。训练时与右移后的目标句算交叉熵。
3.10 训练脚本:手写 warmup 调度器与标签平滑
先补上 2.12 节承诺的两个训练组件:
class NoamScheduler: """2.12.4 节 Noam 学习率调度:lr = d^(-0.5) × min(step^(-0.5), step × warmup^(-1.5)) 前 warmup 步线性爬升 → 之后按 1/√step 衰减。"""
def __init__(self, optimizer, d_model: int, warmup_steps: int = 4000): self.optimizer = optimizer self.d_model = d_model self.warmup_steps = warmup_steps self.step_num = 0
def step(self): """每次参数更新后调用一次。""" self.step_num += 1 lr = (self.d_model ** -0.5) * min( self.step_num ** -0.5, # 衰减段 self.step_num * self.warmup_steps ** -1.5, # 预热段 ) for group in self.optimizer.param_groups: group["lr"] = lr
def get_lr(self): return self.optimizer.param_groups[0]["lr"]
def label_smoothing_loss(logits, target, eps: float = 0.1): """2.12.2 节标签平滑交叉熵。 数学: L = (1-ε)·NLL + ε·均匀惩罚。 第一项让模型往正确类别靠,第二项惩罚"把所有概率押在一个类上"。""" log_probs = F.log_softmax(logits, dim=-1) # (B, K) nll = -log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1).mean() # 正确类负对数似然 smooth = -log_probs.mean(dim=-1).mean() # 对均匀分布的惩罚 return (1 - eps) * nll + eps * smooth
# ---- 验证调度器曲线(前 10 步应为上升期)----tmp_model = nn.Linear(4, 2)tmp_opt = torch.optim.Adam(tmp_model.parameters(), lr=0)sched = NoamScheduler(tmp_opt, d_model=32, warmup_steps=50)for s in range(5): sched.step() print(f"step {s+1}: lr = {sched.get_lr():.6f}")# 输出应逐行增大(预热段线性爬升)3.11 演示 1:完整组装 + toy 信号分类训练
把 3.2~3.7 的全部模块组装成 Encoder-only 分类模型,并完整训练:
class SequenceClassifier(nn.Module): """完整组装:嵌入 → 位置编码 → N 层 Encoder → 读出 → 分类头 参数: readout: 'mean'(池化,最常用)| 'cls'(第 0 个 token)| 'last'(最后 token) pos_mode: 'sinusoidal' | 'learnable' | 'none'('none' 供消融实验) block_cls: 默认 EncoderBlock,消融实验可换成无残差变体 """
def __init__(self, n_layers: int = 2, d_model: int = 32, n_heads: int = 4, num_classes: int = 3, max_len: int = 64, readout: str = "mean", pos_mode: str = "sinusoidal", dropout: float = 0.1, in_features: int = 1, block_cls: nn.Module = EncoderBlock): super().__init__() self.embed = nn.Linear(in_features, d_model) # 标量采样点 → d 维 self.readout = readout if pos_mode == "sinusoidal": self.pos_enc = PositionalEncoding(d_model, max_len, dropout) elif pos_mode == "learnable": self.pos_enc = LearnablePositionalEncoding(max_len, d_model, dropout) else: # 'none':消融用 self.pos_enc = None self.encoder = TransformerEncoder(n_layers, d_model, n_heads, dropout=dropout, block_cls=block_cls) self.head = nn.Linear(d_model, num_classes)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N) 标量序列 x = self.embed(x.unsqueeze(-1)) # (B,N,1) → (B,N,d) if self.pos_enc is not None: x = self.pos_enc(x) # + 位置信息 x = self.encoder(x) # N 层注意力混合 if self.readout == "cls": x = x[:, 0] elif self.readout == "last": x = x[:, -1] else: x = x.mean(dim=1) # 平均池化读出 return self.head(x) # (B, num_classes)
def make_signal_data(n_per_class: int = 300, length: int = 64, seed: int = 0): """3 类"信号":频率 1/2/3 的正弦波 + 高斯噪声。 与两个项目的 AMC(调制识别)场景同构:类别差异藏在频率(周期结构)里。""" torch.manual_seed(seed) xs, ys = [], [] t = torch.arange(length).float() / length * 2 * math.pi # (length,) 0~2π for cls_id, freq in enumerate([1.0, 2.0, 3.0]): for _ in range(n_per_class): phase = torch.rand(1) * 2 * math.pi s = torch.sin(freq * t + phase) + 0.2 * torch.randn(length) xs.append(s) ys.append(cls_id) xs = torch.stack(xs) # (900, 64) ys = torch.tensor(ys) # (900,) perm = torch.randperm(len(ys)) # 打乱顺序 return xs[perm], ys[perm]
def train_demo(): xs, ys = make_signal_data() n_train = int(len(ys) * 0.8) x_train, y_train = xs[:n_train], ys[:n_train] x_test, y_test = xs[n_train:], ys[n_train:]
model = SequenceClassifier() opt = torch.optim.Adam(model.parameters(), lr=1e-3) sched = NoamScheduler(opt, d_model=32, warmup_steps=100) # warmup + 衰减 n_params = sum(p.numel() for p in model.parameters()) print(f"模型参数量: {n_params}")
for step in range(300): # ---- 随机采样一个 batch(玩具数据集不写 DataLoader,专注训练循环本身)---- idx = torch.randint(0, len(x_train), (64,)) xb, yb = x_train[idx], y_train[idx]
opt.zero_grad() # ③ 清空旧梯度 logits = model(xb) # ① 前向 loss = label_smoothing_loss(logits, yb, eps=0.1) # ② 损失(含标签平滑) loss.backward() # ④ 反向 opt.step() # ⑤ 更新 sched.step() # 更新学习率
if step % 50 == 0 or step == 299: model.eval() with torch.no_grad(): acc = (model(x_test).argmax(1) == y_test).float().mean() model.train() print(f"step {step:3d} loss={loss.item():.4f} " f"lr={sched.get_lr():.5f} test_acc={acc:.3f}")
if __name__ == "__main__": train_demo()预期输出(数值有随机波动):
模型参数量: 11815step 0 loss=1.1004 lr=0.00100 test_acc=0.344step 50 loss=0.1268 lr=0.00250 test_acc=0.944step 100 loss=0.0712 lr=0.00354 test_acc=0.989...step 299 loss=0.0385 lr=0.00194 test_acc=1.000观察三个现象:① lr 先爬升后衰减(warmup 生效);② 准确率从 0.34(≈随机)到 1.0;③ 频率不同的正弦波对”无位置信息”的模型是不可分的——下一节的消融实验会证明位置编码在其中的决定性作用。
读到这里你已经完成了”从零实现 + 训练”的全流程。下面两节回答两个问题:模型到底”看”到了什么(3.12 可视化),以及我的实现和官方对不对得上(3.13)。
3.12 演示 2:注意力可视化 + 消融实验
3.12.1 注意力矩阵长什么样
def visualize_attention(): xs, ys = make_signal_data() model = SequenceClassifier()
x = xs[:1] # 取 1 个样本 (1, 64) emb = model.embed(x.unsqueeze(-1)) # (1, 64, 32) emb = model.pos_enc(emb) # + 位置编码 normed = model.encoder.layers[0].norm1(emb) # 第 1 层第 1 个 LN _, attn = model.encoder.layers[0].attn(normed, return_attn=True) w = attn[0, 0] # (64, 64) 第 0 个头
print("注意力矩阵(第 1 层第 1 头,前 8×8 子块,行=query 列=key):") for i in range(8): print(" " + " ".join(f"{w[i, j]:.2f}" for j in range(8)))
# 每行熵:熵小 = 注意力集中(熵 = 分布的不确定性度量:权重均匀分布时熵最大,集中在少数位置时熵趋近 0) row_entropy = -(w * (w + 1e-9).log()).sum(-1) print(f"行熵均值: {row_entropy.mean():.3f}(越小 = 注意力越集中)") print(f"每行权重和: {w.sum(-1)[0]:.3f}(应为 1.0)")运行 visualize_attention()(用 3.11 训练好的模型时注意力更结构化;随机初始化时也能看到行和为 1、权重有差异)。
3.12.2 消融实验:验证每个部件的作用
科学研究的核心方法——逐个拆掉零件,看准确率掉多少:
class NoResidualBlock(EncoderBlock): """消融变体:去掉两条残差连接(对照 2.9.1 节的"梯度高速公路")。"""
def forward(self, x, mask=None, return_attn=False): x = self.attn(self.norm1(x), mask) # 注意:没有 + x x = self.ffn(self.norm2(x)) # 注意:没有 + x return x
def ablation(): xs, ys = make_signal_data() n_train = int(len(ys) * 0.8) x_train, y_train = xs[:n_train], ys[:n_train] x_test, y_test = xs[n_train:], ys[n_train:]
def run(model, name, steps=300): torch.manual_seed(0) # 每个配置同一起跑线 opt = torch.optim.Adam(model.parameters(), lr=1e-3) for step in range(steps): idx = torch.randint(0, len(x_train), (64,)) opt.zero_grad() loss = label_smoothing_loss(model(x_train[idx]), y_train[idx]) loss.backward() opt.step() model.eval() with torch.no_grad(): acc = (model(x_test).argmax(1) == y_test).float().mean() print(f"{name:28s} 最终 test_acc = {acc:.3f}") return acc
run(SequenceClassifier(), "② 完整模型(基准)") run(SequenceClassifier(pos_mode="none"), "① 去掉位置编码") run(SequenceClassifier(block_cls=NoResidualBlock), "③ 去掉残差连接") run(SequenceClassifier(readout="cls"), "④ 换 cls 读出")预期结果与解读:
② 完整模型(基准) 最终 test_acc = 1.000① 去掉位置编码 最终 test_acc ≈ 0.333 ← 崩回随机水平!③ 去掉残差连接 最终 test_acc ≈ 0.3~0.6 ← 训练困难甚至 loss=nan④ 换 cls 读出 最终 test_acc ≈ 1.000 ← 读出方式不改变上限逐条解读:
- ① 去掉位置编码 → 准确率崩回随机(0.333):三类信号的均值、方差完全相同(同为振幅 1 的正弦+同方差噪声),唯一区别是频率——即”随时间变化的模式”。没有位置信息,注意力对顺序置换不变(2.7 节),模型根本”看不见”频率。这个实验把”注意力是置换不变的”从抽象性质变成了亲眼所见的数字;
- ③ 去掉残差 → 无法训练:2 层网络梯度逐层相乘,没有直通路就衰减殆尽(2.9.1 节);
- ④ 换读出方式基本无损:只要 Encoder 学得好,读出手法(cls/mean)影响不大——这也解释了为什么两个项目里各家读出方式(cls、last、GAP、flatten)都能工作。
附加实验(可自行完成):把
d_model从 32 改成 256 观察过拟合(训练 acc 1.0、测试 acc 下降);把n_layers改成 8 观察训练变慢;把学习率改成 1.0 观察 loss 立刻 NaN(对应 5.2 节)。
3.13 演示 3:与官方实现逐层对照(证明自实现正确)
把自实现 EncoderBlock 的每一组权重复制进 nn.TransformerEncoderLayer,若输出逐元素相等,则证明两者数学完全等价:
def compare_with_official(): torch.manual_seed(0) d, H = 32, 4 mine = EncoderBlock(d, H).eval() # 自实现 official = nn.TransformerEncoderLayer( d, H, dim_feedforward=4 * d, batch_first=True, # 输入 (B,N,d) norm_first=True, # Pre-LN activation="gelu", # 与自实现一致 ).eval()
# ---- 权重逐一复制:自实现 qkv(3d,d) ↔ 官方 in_proj_weight(3d,d) ---- official.self_attn.in_proj_weight.data.copy_(mine.attn.qkv.weight.data) official.self_attn.in_proj_bias.data.copy_(mine.attn.qkv.bias.data) official.self_attn.out_proj.weight.data.copy_(mine.attn.proj.weight.data) official.self_attn.out_proj.bias.data.copy_(mine.attn.proj.bias.data) official.linear1.weight.data.copy_(mine.ffn.fc1.weight.data) official.linear1.bias.data.copy_(mine.ffn.fc1.bias.data) official.linear2.weight.data.copy_(mine.ffn.fc2.weight.data) official.linear2.bias.data.copy_(mine.ffn.fc2.bias.data) official.norm1.weight.data.copy_(mine.norm1.weight.data) official.norm1.bias.data.copy_(mine.norm1.bias.data) official.norm2.weight.data.copy_(mine.norm2.weight.data) official.norm2.bias.data.copy_(mine.norm2.bias.data)
x = torch.randn(2, 10, 32) with torch.no_grad(): a, b = mine(x), official(x) print("自实现与官方输出最大误差:", (a - b).abs().max().item()) # 输出 ~1e-7 量级 → 数学完全等价(仅浮点运算顺序差异)
# ---- 三个演示的调用入口(把 3.1~3.13 所有代码块按顺序拼成一个文件时会依次自动执行)----visualize_attention() # 演示 2a:注意力矩阵长什么样(行和为 1、注意力有差异)ablation() # 演示 2b:拆掉位置编码 / 残差 / 换读出,看每个部件的作用compare_with_official() # 演示 3:与官方实现逐层对照,最大误差 ~1e-7这组对应关系值得背下来(复现论文、迁移权重、调参都要用):
| 自实现 | 官方 TransformerEncoderLayer |
|---|---|
attn.qkv.weight (3d, d) | self_attn.in_proj_weight (3d, d) |
attn.proj.weight | self_attn.out_proj.weight |
ffn.fc1 / fc2 | linear1 / linear2 |
norm1 / norm2 | norm1 / norm2 |
| Pre-LN(硬编码) | norm_first=True |
| GELU(硬编码) | activation="gelu" |
3.14 形状推演总表
以 B=64, N=64, d=32, H=4, d_h=8, K=3(演示 1 配置)为例,把整个前向流程过一遍:
| 步骤 | 操作 | 输出形状 |
|---|---|---|
| 1 | 输入信号 | (64, 64) |
| 2 | unsqueeze(-1) | (64, 64, 1) |
| 3 | 线性嵌入 | (64, 64, 32) |
| 4 | + 位置编码(广播) | (64, 64, 32) |
| 5 | LayerNorm1 | (64, 64, 32) |
| 6 | qkv 线性层 | (64, 64, 96) |
| 7 | reshape (B,N,3,H,d_h) | (64, 64, 3, 4, 8) |
| 8 | permute (2,0,3,1,4) | (3, 64, 4, 64, 8) |
| 9 | q / k / v 拆包 | 各 (64, 4, 64, 8) |
| 10 | q @ k^T | (64, 4, 64, 64) |
| 11 | × d_h^(-0.5) | (64, 4, 64, 64) |
| 12 | softmax(dim=-1) | (64, 4, 64, 64)(每行和=1) |
| 13 | attn @ v | (64, 4, 64, 8) |
| 14 | transpose+reshape | (64, 64, 32) |
| 15 | proj(W_O) | (64, 64, 32) |
| 16 | + 残差 x | (64, 64, 32) |
| 17 | LayerNorm2 → FFN → + 残差 | (64, 64, 32) |
| 18 | 第 2 层重复 5~17 | (64, 64, 32) |
| 19 | 末尾 LN | (64, 64, 32) |
| 20 | mean 读出(dim=1) | (64, 32) |
| 21 | 分类头 | (64, 3) |
背下这张表 = 背下整个 Encoder。任何一行形状对不上,报错位置必然在上一行。
第 4 章 框架篇:torch.nn 官方实现
自实现是为了理解;工程和研究里通常直接用官方组件(两个项目亦然:vit.py 自实现做研究,MCDformer 用官方风格组件)。官方接口的参数含义一一对应前文原理。
4.1 nn.MultiheadAttention:官方注意力
attn = nn.MultiheadAttention( embed_dim=32, # = 本文的 d_model num_heads=4, # = n_heads,要求 embed_dim % num_heads == 0 dropout=0.1, # 注意力权重 dropout bias=True, # 投影层是否带偏置 batch_first=True, # ★ True: 输入 (B,N,d);False(默认): (N,B,d) kdim=None, vdim=None, # 交叉注意力用:K/V 的维度(≠embed_dim 时) need_weights=True, # 是否返回注意力权重)
x = torch.randn(2, 10, 32)out, weights = attn(x, x, x) # (query, key, value) —— 自注意力时三者相同print(out.shape) # (2, 10, 32)print(weights.shape) # (2, 10, 10):平均了所有头的权重
# 两种掩码(语义不同,别混!)key_pad = torch.zeros(2, 10, dtype=torch.bool); key_pad[:, 0] = True # (B,N) True=屏蔽att_mask = torch.zeros(10, 10, dtype=torch.bool) # (N,N) True=屏蔽out2, _ = attn(x, x, x, key_padding_mask=key_pad, attn_mask=att_mask)参数对照自实现:
embed_dim、num_heads↔d_model、n_heads;- 官方用三个独立投影层(timm 风格),合并写法
in_proj_weight即本文的qkv; key_padding_mask(形状 (B,N),True=屏蔽)与attn_mask(形状 (N,N) 或 (N·H,N,N),True=屏蔽)语义与自实现的mask==0 屏蔽相反——官方用 True 屏蔽;- 交叉注意力用法:
attn(query=decoder_x, key=enc_out, value=enc_out),Q 与 K/V 可以来自不同序列——这正是 3.8 节CrossAttention的官方版。
4.2 nn.TransformerEncoderLayer / Encoder:官方编码块
encoder_layer = nn.TransformerEncoderLayer( d_model=32, nhead=4, dim_feedforward=128, # = 4*d_model,本文的 d_ff dropout=0.1, activation="gelu", # 或 "relu"(默认) batch_first=True, # ★ (B,N,d) norm_first=True, # ★ True=Pre-LN;False=Post-LN(原论文) layer_norm_eps=1e-5, # LN 分母里的 ε(2.9.2 节))encoder = nn.TransformerEncoder(encoder_layer, num_layers=2)out = encoder(torch.randn(2, 10, 32)) # (2, 10, 32)官方 Encoder 不内置末尾 LN(本文 3.7 节的 self.norm 是 ViT 惯例的额外添加),复现时注意这层差异。
4.3 nn.Transformer:官方完整 Encoder-Decoder + 翻译玩具
transformer = nn.Transformer( d_model=32, nhead=4, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=128, dropout=0.1, batch_first=True, # 输入输出 (B,N,d))
src = torch.randint(0, 50, (2, 8)) # 源句词编号tgt = torch.randint(0, 50, (2, 5)) # 目标句(右移一位)src_emb = nn.Embedding(50, 32)(src) # 官方不含嵌入层,需自备tgt_emb = nn.Embedding(50, 32)(tgt)
tgt_mask = nn.Transformer.generate_square_subsequent_mask(5) # 因果掩码 (5,5)out = transformer(src_emb, tgt_emb, tgt_mask=tgt_mask)print(out.shape) # (2, 5, 32) —— 每个位置一个 d 维表示注意:官方 nn.Transformer 是”裸结构”——不含词嵌入、位置编码、输出 softmax 层,这些都要像 3.9 节那样自己接。
4.4 自实现 ↔ 官方 完整对照表
| 本文自实现 | 官方组件 | 差异说明 |
|---|---|---|
MultiHeadAttention | nn.MultiheadAttention | 官方多 kdim/vdim(交叉)、key_padding_mask |
PositionwiseFeedForward | linear1→act→dropout→linear2(内嵌) | 官方内嵌在 Layer 里 |
EncoderBlock | nn.TransformerEncoderLayer(norm_first=True) | 官方默认 Post-LN,必须改参数 |
TransformerEncoder | nn.TransformerEncoder | 官方无末尾 LN |
DecoderBlock | nn.TransformerDecoderLayer | 官方参数更全(含 norm_first) |
Transformer | nn.Transformer | 官方不含嵌入/位置编码/head |
PositionalEncoding | 无(需自写) | 官方从未内置位置编码 |
NoamScheduler | torch.optim.lr_scheduler.LambdaLR | 官方可用 lambda 一行实现 |
4.5 与本仓库两个项目的代码对照
| 项目文件 | 实现风格 | 与本文的对应 |
|---|---|---|
vit.py MultiHeadAttention | 合并 qkv + permute(2,0,3,1,4) | 与 3.4 节逐行一致 |
vit.py ViTBlock | Pre-LN | 与 3.6 节一致(仅 FFN 顺序微差) |
MCDformer.py Attention | q/k/v 三个独立线性层(timm 风格) | 3.4 节注释所述变体,数学等价 |
MCDformer.py Block | Pre-LN + GELU + drop_path | 与 3.6 节一致,dropout 强度 0.5 |
CTDNN.py Block | x = norm1(x + attn(x)) | Post-LN 变体(特化文档第 5 章分析) |
第 5 章 调试篇:报错与陷阱
每一条都是真实会发生的错误。报错时从错误信息最后一行往上读,找到你自己代码的行号,再对照本表。
5.1 常见报错速查表
5.1.1 矩阵乘法维度不匹配
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x64 and 1x32)解读:(2×64) @ (1×32) —— 左矩阵列数 64 ≠ 右矩阵行数 1。你的张量形状与 nn.Linear(in_features, ...) 声明不一致。常见原因:① 忘了 unsqueeze((B,N) 直接送进 Linear(1,d));② 序列维和特征维写反((B,d,N) 送进 Linear(d,…))。
修复:在出错行前面打印 print(x.shape),对照 3.14 节总表逐行核对。
5.1.2 reshape 后形状对不上
RuntimeError: shape '[2, 10, 3, 4, 8]' is invalid for input of size 2560解读:2×10×3×4×8 = 1920 ≠ 2560。你的 reshape 目标形状和实际元素总数不符——多半是忘了 permute((B,N,H,d_h) 直接 reshape 成 (B,N,d) 之类)。
修复:reshape 之前先 print(x.shape);牢记”先 permute 换轴、再 reshape 合并相邻轴”。
5.1.3 索引维数不对
IndexError: too many indices for tensor of dimension 3解读:对 3 维张量用了 4 个下标(如 x[:, :, 0, :])。常见原因:数据集产出 (B,2,L),模型却按 (B,1,2,L) 写 forward——数据集与模型之间的”通道契约”不一致(本仓库 meta-transformer 的元学习路径就存在这个隐患,详见特化文档 7.4 节)。
修复:统一在模型入口 x = x.unsqueeze(1),或数据集出口 np.expand_dims(x, axis=1),二选一并写断言。
5.1.4 batch_first 混用(官方 API 专属)
RuntimeError: shape '[2, 10, 32]' is invalid for input of size 640解读:nn.MultiheadAttention 默认 batch_first=False,期望输入 (N,B,d),你给了 (B,N,d)(或反之)。
修复:创建层时显式写 batch_first=True,全项目统一。这是官方 API 第一大坑。
5.1.5 掩码数据类型错误
RuntimeError: masked_fill__(): value with type Float cannot be cast to type Bool解读:masked_fill(mask == 0, -inf) 要求 mask 是 bool 张量。你把浮点掩码直接传进来了。
修复:mask = mask.bool(),或比较生成 bool:(pad_mask == 0)。
5.1.6 inplace 修改破坏梯度
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation解读:对参与计算图的张量做了原地修改(x += ...、x[0] = ...)。
修复:x = x + ...(新建张量)代替 x += ...;给不需要梯度的操作加 with torch.no_grad():。
5.1.7 显存不足
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB ...解读:注意力矩阵是 (5.3 节)。N=4096、batch=32 时单层注意力图约 ≈ 数 GB。
修复:减小 batch、减小序列长度、用 torch.utils.checkpoint、或改稀疏/分块注意力。
5.2 NaN 排查决策树
loss 变成 NaN 后按顺序问:
loss = NaN ├─ 学习率是否 > 1e-2 且无 warmup? → 是:降到 1e-4 并加 NoamScheduler ├─ 是否删了 √d_h 缩放且 d_h 很大? → 是:补上缩放(2.4 节) ├─ mask 是否把某行整个屏蔽了? → 是:softmax 全 -inf → NaN,检查 padding ├─ 输入数据是否含 inf/超大值? → 是:检查数据归一化 ├─ 是否用了 fp16 且没 GradScaler? → 是:见 5.4 节 └─ 是否改过残差/LN 顺序? → 是:对照 2.9.3 恢复 Pre-LN验证手法:在可疑位置插入
assert torch.isfinite(x).all(), f"出现 NaN/Inf,形状 {x.shape}"哪个断言先炸,问题就在它前面。
5.3 复杂度与显存(为什么序列长度是生命线)
自注意力:时间 、显存 (注意力矩阵 ,每头每样本一张)。
| 序列长度 N | 单头单样本注意力矩阵 | 单层 8 头 batch=32 |
|---|---|---|
| 128 | 64 KB | 16 MB |
| 1024 | 4 MB | 1 GB |
| 8192 | 256 MB | 64 GB(爆炸) |
两个项目正是吃准了这一点:信号长度 1281024 恰好落在自注意力最擅长的区间(这也是特化文档反复强调”token 数控制在 16512”的原因)。超长序列请转向分块注意力(Swin)、稀疏注意力(Longformer)、线性注意力(Linformer)或 FlashAttention。
5.4 混合精度(AMP)注意
fp16(16 位半精度浮点,占用显存减半、速度更快,但数值表示范围小)训练可提速省显存,但 softmax 与累加对精度敏感:
scaler = torch.cuda.amp.GradScaler() # 自动缩放梯度防下溢with torch.cuda.amp.autocast(): logits = model(x) # 前向自动混精度 loss = loss_fn(logits, y)scaler.scale(loss).backward() # 注意:不是 loss.backward()scaler.step(opt)scaler.update()若仍 NaN:把注意力打分与 softmax 强制回 fp32(attn.float().softmax(-1)),或先不用 AMP 跑通再开。
5.5 最终检查清单(提交/复现前过一遍)
-
d_model % n_heads == 0 - 位置编码只加一次、加在第一层之前
- 残差两条支路都在(
x = x + ...) - Pre-LN 与 Post-LN 全项目统一(推荐 Pre-LN)
- 官方 API 显式
batch_first=True - mask 语义统一(本文自实现 0=屏蔽;官方 True=屏蔽)
- 训练有 warmup;lr ≤ 1e-3
- 数据集出口形状与模型入口形状一致(写断言)
- 评估时
model.eval()(否则 dropout 干扰结果)
第 6 章 小结与下一步
6.1 本文知识地图
数学准备(第 0 章)── 点积=相似度、矩阵乘=加权组合、softmax=归一化 │动机(第 1 章)────── RNN 三大病 → 注意力"圆桌会议" │原理(第 2 章)────── 手算注意力 → 缩放/多头/位置/LN/残差/掩码/训练三件套 │代码(第 3 章)────── NumPy 验证数学 → 8 个模块逐行 → 训练 → 可视化 → 消融 → 对照官方 │框架(第 4 章)────── nn.MultiheadAttention / TransformerEncoderLayer / Transformer │调试(第 5 章)────── 7 类报错 + NaN 决策树 + O(N²) + AMP + 检查清单6.2 五个”一句话记住”
- 注意力 = ——“打分、归一化、加权搬运”;
- 多头 = 把 d 拆 H 份并行开会,最后拼接投影;
- 位置编码解决注意力看不见顺序的问题;
- 残差 + Pre-LN 是深网络可训练的命根子;
- 训练三件套:交叉熵(+标签平滑)、Adam、warmup。
6.3 下一步
掌握这套”标准件”后,真正的研究从”改造”开始:针对自己的任务,把 token 化方式、注意力作用域、位置编码、块结构、任务头、训练范式逐一重设计——这就是本仓库两个项目做的事(Meta-Transformer 把 ViT 改造成小样本元学习器,MCDformer 把注意力搬到通道轴并前置频域去噪模块)。完整方法论与逐行拆解请读姊妹文档:
附录:完整可运行代码(第 3 章全部模块 + 三个演示,已在 GPU 上验证)
把下面代码整体复制保存为
transformer_from_scratch.py,运行python transformer_from_scratch.py。 默认在 GPU 上运行,全程约 13 分钟。 按顺序自动执行:3.1 NumPy 注意力手算复现 → 3.23.10 各模块冒烟测试 → 演示 1(toy 信号分类训练,test_acc 应达 0.95 以上)→ 演示 2(注意力可视化 + 消融实验)→ 演示 3(与官方实现逐层对照,最大误差约 1e-7)。 与正文 3.1~3.13 各代码块内容一致,只是合并为单一文件并修复了两处拼装问题(InputEmbedding保存d_model、掩码广播形状)。
import numpy as np
def softmax(z, axis=-1): """数值稳定版 softmax(2.5 节):先减最大值再指数。""" z = z - z.max(axis=axis, keepdims=True) # 防止 exp 溢出 e = np.exp(z) return e / e.sum(axis=axis, keepdims=True)
def attention_numpy(X, Wq, Wk, Wv, d_k): """X: (N, d) 一个序列。返回注意力输出与权重矩阵。""" Q = X @ Wq # (N, d_k) —— 投影出查询 K = X @ Wk # (N, d_k) —— 投影出键 V = X @ Wv # (N, d_k) —— 投影出值 scores = Q @ K.T / np.sqrt(d_k) # (N, N) —— 相似度打分并缩放 A = softmax(scores, axis=1) # (N, N) —— 逐行 softmax out = A @ V # (N, d_k) —— 加权组合 return out, A
# ---- 复现 2.2 节的例子 ----X = np.array([[1.0, 0.0], [0.0, 1.0]]) # 2 个 token, 每 token 2 维I = np.eye(2) # 单位矩阵: Q=K=V=Xout, A = attention_numpy(X, I, I, I, d_k=2)
print("注意力权重 A =\n", np.round(A, 3))print("输出 =\n", np.round(out, 3))
import mathimport torchimport torch.nn as nnimport torch.nn.functional as F
torch.manual_seed(0) # 固定随机种子,保证结果可复现
# ---- 设备选择:默认 GPU ----# set_default_device 让之后创建的所有张量/模块自动落在 GPU 上,正文代码无需再写 .to()/.cuda()torch.set_default_device('cuda')print("默认设备: cuda ——", torch.cuda.get_device_name(0))
class InputEmbedding(nn.Module): """两种嵌入方式: vocab 不为 None: 查表嵌入(NLP 用)——token 是词的编号; vocab 为 None : 线性嵌入(信号/数值用)——token 是连续数值。 """
def __init__(self, d_model: int, vocab: int = None, in_features: int = 1): super().__init__() self.d_model = d_model if vocab is not None: self.embed = nn.Embedding(vocab, d_model) # 查表: (B,N) → (B,N,d) else: self.embed = nn.Linear(in_features, d_model) # 线性: (B,N,1) → (B,N,d)
def forward(self, x: torch.Tensor) -> torch.Tensor: return self.embed(x) * math.sqrt(self.d_model)
# ---- 使用示例 ----emb = InputEmbedding(d_model=16, vocab=1000) # NLP 风格ids = torch.randint(0, 1000, (2, 5)) # (B=2, N=5) 词编号print(emb(ids).shape) # (2, 5, 16)
emb2 = InputEmbedding(d_model=16, in_features=1) # 信号风格sig = torch.randn(2, 5, 1) # (B=2, N=5, 1) 标量序列print(emb2(sig).shape) # (2, 5, 16)
class PositionalEncoding(nn.Module): """正弦位置编码(2.7.1 节)。register_buffer 的内容不参与训练但会随模型保存。"""
def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1): super().__init__() pe = torch.zeros(max_len, d_model) # (max_len, d) 先全 0 position = torch.arange(0, max_len).unsqueeze(1).float() # (max_len, 1): 0,1,2,... # 分母 10000^(2i/d):用 exp(log) 形式写,数值更稳(对 i=0,2,4,... 生成 d/2 个) div_term = torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) # (d/2,) pe[:, 0::2] = torch.sin(position * div_term) # 偶数维填 sin pe[:, 1::2] = torch.cos(position * div_term) # 奇数维填 cos self.register_buffer("pe", pe.unsqueeze(0)) # (1, max_len, d) 预备广播 self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N, d)。pe[:, :N] 取前 N 个位置,靠广播加到每个样本上 x = x + self.pe[:, : x.size(1), :] return self.dropout(x)
class LearnablePositionalEncoding(nn.Module): """可学习位置编码(2.7.2 节,ViT 与两个项目同款)。"""
def __init__(self, num_tokens: int, d_model: int, dropout: float = 0.1): super().__init__() self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, d_model)) self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, num_tokens, d) return self.dropout(x + self.pos_embed)
# ---- 验证:位置编码形状与数值 ----pe = PositionalEncoding(d_model=4, max_len=100)x = torch.zeros(2, 3, 4) # (B=2, N=3, d=4)out = pe(x)print(out.shape) # (2, 3, 4) —— 形状不变,只加信息print("位置 0 编码:", pe.pe[0, 0].tolist()) # [0, 1, 0, 1] —— 对照 2.7.1 手算print("位置 1 编码:", [round(v, 3) for v in pe.pe[0, 1].tolist()]) # ≈ [0.841, 0.540, 0.010, 1.0]
class MultiHeadAttention(nn.Module): """多头自注意力。实现方式:Q/K/V 合并为一个大线性层(原论文同款)。
形状流程(对照 2.6 节表格): x (B,N,d) → qkv (B,N,3d) → reshape (B,N,3,H,d_h) → permute (3,B,H,N,d_h) → scores (B,H,N,N) → softmax → attn@V → (B,H,N,d_h) → transpose+reshape (B,N,d) → proj (B,N,d) """
def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0, f"d_model({d_model}) 必须能被 n_heads({n_heads}) 整除" self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads # 每个头的维度 d_h self.qkv = nn.Linear(d_model, 3 * d_model) # 一次投影出 Q,K,V(省 3 次调用) self.proj = nn.Linear(d_model, d_model) # 输出投影 W_O self.attn_drop = nn.Dropout(dropout) self.proj_drop = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): B, N, d = x.shape
# ---- 第 1 步:投影 + 拆出 Q/K/V ---- # Linear 输出 (B,N,3d);reshape 成 (B,N,3,H,d_h); # permute(2,0,3,1,4) 把"第 2 维(3)"提到最前 → (3,B,H,N,d_h),方便拆包 qkv = ( self.qkv(x) .reshape(B, N, 3, self.n_heads, self.d_head) .permute(2, 0, 3, 1, 4) ) q, k, v = qkv[0], qkv[1], qkv[2] # 各 (B,H,N,d_h)
# ---- 第 2 步:打分 + 缩放 ---- # q @ k^T: (B,H,N,d_h)×(B,H,N,d_h)^T → (B,H,N,N) # 语义:每个头里,每个 query 与每个 key 的相似度 attn = (q @ k.transpose(-2, -1)) * (self.d_head ** -0.5)
# ---- 第 3 步:掩码(可选)---- # mask 为 0 的位置置 -inf;softmax(exp(-inf)=0) 后权重为 0 # 支持两种形状:(B,N) padding 掩码 → (B,1,1,N);(N,N) 因果掩码 → (1,1,N,N) if mask is not None: if mask.dim() == 2 and mask.shape == attn.shape[-2:]: mask = mask.unsqueeze(0).unsqueeze(0) # 方阵掩码:广播到 batch 与头 else: mask = mask.unsqueeze(1).unsqueeze(2) # 逐样本掩码:广播到头与 query attn = attn.masked_fill(mask == 0, float("-inf"))
# ---- 第 4 步:softmax 归一化(沿最后一维 = 对 key 求和为 1)---- attn = attn.softmax(dim=-1) attn = self.attn_drop(attn)
# ---- 第 5 步:加权求和 + 拼头 + 输出投影 ---- out = attn @ v # (B,H,N,d_h) out = out.transpose(1, 2).reshape(B, N, d) # (B,N,d):把头拼回特征维 out = self.proj_drop(self.proj(out))
if return_attn: return out, attn # 附带注意力图(可视化用) return out
mha = MultiHeadAttention(d_model=32, n_heads=4)x = torch.randn(2, 10, 32) # (B=2, N=10, d=32)out, attn = mha(x, return_attn=True)print("输出形状:", out.shape) # (2, 10, 32) —— 形状不变print("注意力图形状:", attn.shape) # (2, 4, 10, 10) —— 每头一张 10×10 图print("每行权重和:", attn[0, 0, 0].sum().item()) # ≈1.0 —— softmax 性质
# 可训练性检查:参数能收到梯度loss = out.sum()loss.backward()print("qkv 权重梯度非空:", mha.qkv.weight.grad is not None) # True
# 掩码检查:屏蔽第 0 个 key 后,第 0 列权重应为 0mask = torch.ones(2, 10); mask[:, 0] = 0 # (B,N) 形状的 padding 掩码out_m, attn_m = mha(x, mask=mask, return_attn=True)print("被屏蔽列权重:", attn_m[0, 0, :, 0].sum().item()) # 0.0 ✓
class PositionwiseFeedForward(nn.Module): """FFN(2.8 节): d → 4d → d。逐 token 独立,形状永不改变。"""
def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1): super().__init__() d_ff = d_ff or 4 * d_model self.fc1 = nn.Linear(d_model, d_ff) # 放大 4 倍给"加工空间" self.fc2 = nn.Linear(d_ff, d_model) # 缩回原维度 self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N, d) —— Linear 自动作用在最后一维,对每个 token 都一样 return self.fc2(self.dropout(F.gelu(self.fc1(x))))
ffn = PositionwiseFeedForward(32)print(ffn(torch.randn(2, 10, 32)).shape) # (2, 10, 32)
class EncoderBlock(nn.Module): """Pre-LN 编码块(2.9.3 节): x = x + Attn(LN(x)) ← 先归一化再进子层,最后残差相加 x = x + FFN(LN(x)) """
def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1): super().__init__() self.norm1 = nn.LayerNorm(d_model) # 注意力前的 LN self.attn = MultiHeadAttention(d_model, n_heads, dropout) self.norm2 = nn.LayerNorm(d_model) # FFN 前的 LN self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): if return_attn: a_out, attn = self.attn(self.norm1(x), mask, return_attn=True) x = x + a_out # 残差:梯度直通 else: x = x + self.attn(self.norm1(x), mask) x = x + self.ffn(self.norm2(x)) # FFN 支路 + 残差 if return_attn: return x, attn return x
block = EncoderBlock(32, 4)x = torch.randn(2, 10, 32)out, attn = block(x, return_attn=True)print("块输出形状:", out.shape) # (2, 10, 32)print("块内注意力图:", attn.shape) # (2, 4, 10, 10)
class TransformerEncoder(nn.Module): def __init__(self, n_layers: int, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1, block_cls: nn.Module = EncoderBlock): super().__init__() # ModuleList 才能被优化器"看见"全部参数;普通 list 不会注册参数! # block_cls 参数让消融实验可以注入无残差等变体块 self.layers = nn.ModuleList( [block_cls(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)] ) self.norm = nn.LayerNorm(d_model) # 末尾 LN(ViT 惯例)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, return_attn: bool = False): all_attns = [] for layer in self.layers: if return_attn: x, attn = layer(x, mask, return_attn=True) all_attns.append(attn) else: x = layer(x, mask) if return_attn: return self.norm(x), all_attns return self.norm(x)
encoder = TransformerEncoder(n_layers=2, d_model=32, n_heads=4)x = torch.randn(2, 10, 32)out, attns = encoder(x, return_attn=True)print("编码器输出:", out.shape) # (2, 10, 32)print("层数 × 每层注意力图:", len(attns), attns[0].shape) # 2 张 (2,4,10,10)
class DecoderBlock(nn.Module): """解码块(2.11 节)。与编码块的区别: ① 自注意力加因果掩码(只看过去); ② 多一个交叉注意力(Q 来自解码器,K/V 来自编码器输出)。"""
def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1): super().__init__() self.norm1 = nn.LayerNorm(d_model) self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) # ① 掩码自注意力 self.norm2 = nn.LayerNorm(d_model) self.cross_attn = CrossAttention(d_model, n_heads, dropout) # ② 交叉注意力 self.norm3 = nn.LayerNorm(d_model) self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
def forward(self, x: torch.Tensor, enc_out: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor: x = x + self.self_attn(self.norm1(x), causal_mask) # 只许看 j ≤ i x = x + self.cross_attn(self.norm2(x), enc_out) # 向编码器"提问" x = x + self.ffn(self.norm3(x)) return x
class CrossAttention(nn.Module): """交叉注意力:Q 来自 x(解码端),K/V 来自 enc(编码端)。 与自注意力的唯一区别:K/V 的投影层接收的是另一个张量。"""
def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0 self.d_model, self.n_heads = d_model, n_heads self.d_head = d_model // n_heads self.wq = nn.Linear(d_model, d_model) # Q 投影(用于解码端) self.wk = nn.Linear(d_model, d_model) # K 投影(用于编码端) self.wv = nn.Linear(d_model, d_model) # V 投影(用于编码端) self.proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, enc: torch.Tensor) -> torch.Tensor: B, N, d = x.shape E = enc.size(1) # 编码端 token 数(可与 N 不同) q = self.wq(x).reshape(B, N, self.n_heads, self.d_head).permute(0, 2, 1, 3) k = self.wk(enc).reshape(B, E, self.n_heads, self.d_head).permute(0, 2, 1, 3) v = self.wv(enc).reshape(B, E, self.n_heads, self.d_head).permute(0, 2, 1, 3) attn = (q @ k.transpose(-2, -1)) * (self.d_head ** -0.5) # (B,H,N,E) 注意是 N×E! attn = attn.softmax(dim=-1) out = (attn @ v).transpose(1, 2).reshape(B, N, d) return self.proj(self.dropout(out))
def make_causal_mask(n: int, device=None) -> torch.Tensor: """下三角矩阵(2.10 节):(i,j) 处 j ≤ i 才为 1。""" return torch.tril(torch.ones(n, n, device=device))
# ---- 验证 ----cm = make_causal_mask(4)print(cm) # 对照 2.10 节的 4×4 矩阵dec_block = DecoderBlock(32, 4)x = torch.randn(2, 5, 32) # 解码端 5 个 tokenenc = torch.randn(2, 10, 32) # 编码端 10 个 tokenout = dec_block(x, enc, make_causal_mask(5))print("解码块输出:", out.shape) # (2, 5, 32) —— 解码端 token 数不变
class Transformer(nn.Module): """原版 Encoder-Decoder 结构(2.11 节全景图)。"""
def __init__(self, vocab_size: int, d_model: int = 32, n_heads: int = 4, n_enc: int = 2, n_dec: int = 2, d_ff: int = None, dropout: float = 0.1, max_len: int = 100): super().__init__() self.d_model = d_model self.src_embed = nn.Embedding(vocab_size, d_model) # 源词嵌入 self.tgt_embed = nn.Embedding(vocab_size, d_model) # 目标词嵌入 self.pos_enc = PositionalEncoding(d_model, max_len, dropout) # 共用一套位置编码 self.encoder = TransformerEncoder(n_enc, d_model, n_heads, d_ff, dropout) self.decoder = nn.ModuleList( [DecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_dec)] ) self.head = nn.Linear(d_model, vocab_size) # 词表概率
def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor: tgt_len = tgt.size(1) enc_out = self.encoder(self.pos_enc(self.src_embed(src))) # (B,S,d) x = self.pos_enc(self.tgt_embed(tgt)) # (B,T,d) causal = make_causal_mask(tgt_len, tgt.device) # (T,T) for blk in self.decoder: x = blk(x, enc_out, causal) return self.head(x) # (B,T,vocab)
# ---- 形状冒烟测试 ----model = Transformer(vocab_size=100)src = torch.randint(0, 100, (2, 8)) # 源句 8 词tgt = torch.randint(0, 100, (2, 5)) # 目标句 5 词logits = model(src, tgt)print("输出 logits 形状:", logits.shape) # (2, 5, 100) —— 每个位置预测下一个词
class NoamScheduler: """2.12.4 节 Noam 学习率调度:lr = d^(-0.5) × min(step^(-0.5), step × warmup^(-1.5)) 前 warmup 步线性爬升 → 之后按 1/√step 衰减。"""
def __init__(self, optimizer, d_model: int, warmup_steps: int = 4000): self.optimizer = optimizer self.d_model = d_model self.warmup_steps = warmup_steps self.step_num = 0
def step(self): """每次参数更新后调用一次。""" self.step_num += 1 lr = (self.d_model ** -0.5) * min( self.step_num ** -0.5, # 衰减段 self.step_num * self.warmup_steps ** -1.5, # 预热段 ) for group in self.optimizer.param_groups: group["lr"] = lr
def get_lr(self): return self.optimizer.param_groups[0]["lr"]
def label_smoothing_loss(logits, target, eps: float = 0.1): """2.12.2 节标签平滑交叉熵。 数学: L = (1-ε)·NLL + ε·均匀惩罚。 第一项让模型往正确类别靠,第二项惩罚"把所有概率押在一个类上"。""" log_probs = F.log_softmax(logits, dim=-1) # (B, K) nll = -log_probs.gather(-1, target.unsqueeze(-1)).squeeze(-1).mean() # 正确类负对数似然 smooth = -log_probs.mean(dim=-1).mean() # 对均匀分布的惩罚 return (1 - eps) * nll + eps * smooth
# ---- 验证调度器曲线(前 10 步应为上升期)----tmp_model = nn.Linear(4, 2)tmp_opt = torch.optim.Adam(tmp_model.parameters(), lr=0)sched = NoamScheduler(tmp_opt, d_model=32, warmup_steps=50)for s in range(5): sched.step() print(f"step {s+1}: lr = {sched.get_lr():.6f}")# 输出应逐行增大(预热段线性爬升)
class SequenceClassifier(nn.Module): """完整组装:嵌入 → 位置编码 → N 层 Encoder → 读出 → 分类头 参数: readout: 'mean'(池化,最常用)| 'cls'(第 0 个 token)| 'last'(最后 token) pos_mode: 'sinusoidal' | 'learnable' | 'none'('none' 供消融实验) block_cls: 默认 EncoderBlock,消融实验可换成无残差变体 """
def __init__(self, n_layers: int = 2, d_model: int = 32, n_heads: int = 4, num_classes: int = 3, max_len: int = 64, readout: str = "mean", pos_mode: str = "sinusoidal", dropout: float = 0.1, in_features: int = 1, block_cls: nn.Module = EncoderBlock): super().__init__() self.embed = nn.Linear(in_features, d_model) # 标量采样点 → d 维 self.readout = readout if pos_mode == "sinusoidal": self.pos_enc = PositionalEncoding(d_model, max_len, dropout) elif pos_mode == "learnable": self.pos_enc = LearnablePositionalEncoding(max_len, d_model, dropout) else: # 'none':消融用 self.pos_enc = None self.encoder = TransformerEncoder(n_layers, d_model, n_heads, dropout=dropout, block_cls=block_cls) self.head = nn.Linear(d_model, num_classes)
def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (B, N) 标量序列 x = self.embed(x.unsqueeze(-1)) # (B,N,1) → (B,N,d) if self.pos_enc is not None: x = self.pos_enc(x) # + 位置信息 x = self.encoder(x) # N 层注意力混合 if self.readout == "cls": x = x[:, 0] elif self.readout == "last": x = x[:, -1] else: x = x.mean(dim=1) # 平均池化读出 return self.head(x) # (B, num_classes)
def make_signal_data(n_per_class: int = 300, length: int = 64, seed: int = 0): """3 类"信号":频率 1/2/3 的正弦波 + 高斯噪声。 与两个项目的 AMC(调制识别)场景同构:类别差异藏在频率(周期结构)里。""" torch.manual_seed(seed) xs, ys = [], [] t = torch.arange(length).float() / length * 2 * math.pi # (length,) 0~2π for cls_id, freq in enumerate([1.0, 2.0, 3.0]): for _ in range(n_per_class): phase = torch.rand(1) * 2 * math.pi s = torch.sin(freq * t + phase) + 0.2 * torch.randn(length) xs.append(s) ys.append(cls_id) xs = torch.stack(xs) # (900, 64) ys = torch.tensor(ys) # (900,) perm = torch.randperm(len(ys)) # 打乱顺序 return xs[perm], ys[perm]
def train_demo(): xs, ys = make_signal_data() n_train = int(len(ys) * 0.8) x_train, y_train = xs[:n_train], ys[:n_train] x_test, y_test = xs[n_train:], ys[n_train:]
model = SequenceClassifier() opt = torch.optim.Adam(model.parameters(), lr=1e-3) sched = NoamScheduler(opt, d_model=32, warmup_steps=100) # warmup + 衰减 n_params = sum(p.numel() for p in model.parameters()) print(f"模型参数量: {n_params}")
for step in range(300): # ---- 随机采样一个 batch(玩具数据集不写 DataLoader,专注训练循环本身)---- idx = torch.randint(0, len(x_train), (64,)) xb, yb = x_train[idx], y_train[idx]
opt.zero_grad() # ③ 清空旧梯度 logits = model(xb) # ① 前向 loss = label_smoothing_loss(logits, yb, eps=0.1) # ② 损失(含标签平滑) loss.backward() # ④ 反向 opt.step() # ⑤ 更新 sched.step() # 更新学习率
if step % 50 == 0 or step == 299: model.eval() with torch.no_grad(): acc = (model(x_test).argmax(1) == y_test).float().mean() model.train() print(f"step {step:3d} loss={loss.item():.4f} " f"lr={sched.get_lr():.5f} test_acc={acc:.3f}")
if __name__ == "__main__": train_demo()
def visualize_attention(): xs, ys = make_signal_data() model = SequenceClassifier()
x = xs[:1] # 取 1 个样本 (1, 64) emb = model.embed(x.unsqueeze(-1)) # (1, 64, 32) emb = model.pos_enc(emb) # + 位置编码 normed = model.encoder.layers[0].norm1(emb) # 第 1 层第 1 个 LN _, attn = model.encoder.layers[0].attn(normed, return_attn=True) w = attn[0, 0] # (64, 64) 第 0 个头
print("注意力矩阵(第 1 层第 1 头,前 8×8 子块,行=query 列=key):") for i in range(8): print(" " + " ".join(f"{w[i, j]:.2f}" for j in range(8)))
# 每行熵:熵小 = 注意力集中(熵 = 分布的不确定性度量:权重均匀分布时熵最大,集中在少数位置时熵趋近 0) row_entropy = -(w * (w + 1e-9).log()).sum(-1) print(f"行熵均值: {row_entropy.mean():.3f}(越小 = 注意力越集中)") print(f"每行权重和: {w.sum(-1)[0]:.3f}(应为 1.0)")
class NoResidualBlock(EncoderBlock): """消融变体:去掉两条残差连接(对照 2.9.1 节的"梯度高速公路")。"""
def forward(self, x, mask=None, return_attn=False): x = self.attn(self.norm1(x), mask) # 注意:没有 + x x = self.ffn(self.norm2(x)) # 注意:没有 + x return x
def ablation(): xs, ys = make_signal_data() n_train = int(len(ys) * 0.8) x_train, y_train = xs[:n_train], ys[:n_train] x_test, y_test = xs[n_train:], ys[n_train:]
def run(model, name, steps=300): torch.manual_seed(0) # 每个配置同一起跑线 opt = torch.optim.Adam(model.parameters(), lr=1e-3) for step in range(steps): idx = torch.randint(0, len(x_train), (64,)) opt.zero_grad() loss = label_smoothing_loss(model(x_train[idx]), y_train[idx]) loss.backward() opt.step() model.eval() with torch.no_grad(): acc = (model(x_test).argmax(1) == y_test).float().mean() print(f"{name:28s} 最终 test_acc = {acc:.3f}") return acc
run(SequenceClassifier(), "② 完整模型(基准)") run(SequenceClassifier(pos_mode="none"), "① 去掉位置编码") run(SequenceClassifier(block_cls=NoResidualBlock), "③ 去掉残差连接") run(SequenceClassifier(readout="cls"), "④ 换 cls 读出")
def compare_with_official(): torch.manual_seed(0) d, H = 32, 4 mine = EncoderBlock(d, H).eval() # 自实现 official = nn.TransformerEncoderLayer( d, H, dim_feedforward=4 * d, batch_first=True, # 输入 (B,N,d) norm_first=True, # Pre-LN activation="gelu", # 与自实现一致 ).eval()
# ---- 权重逐一复制:自实现 qkv(3d,d) ↔ 官方 in_proj_weight(3d,d) ---- official.self_attn.in_proj_weight.data.copy_(mine.attn.qkv.weight.data) official.self_attn.in_proj_bias.data.copy_(mine.attn.qkv.bias.data) official.self_attn.out_proj.weight.data.copy_(mine.attn.proj.weight.data) official.self_attn.out_proj.bias.data.copy_(mine.attn.proj.bias.data) official.linear1.weight.data.copy_(mine.ffn.fc1.weight.data) official.linear1.bias.data.copy_(mine.ffn.fc1.bias.data) official.linear2.weight.data.copy_(mine.ffn.fc2.weight.data) official.linear2.bias.data.copy_(mine.ffn.fc2.bias.data) official.norm1.weight.data.copy_(mine.norm1.weight.data) official.norm1.bias.data.copy_(mine.norm1.bias.data) official.norm2.weight.data.copy_(mine.norm2.weight.data) official.norm2.bias.data.copy_(mine.norm2.bias.data)
x = torch.randn(2, 10, 32) with torch.no_grad(): a, b = mine(x), official(x) print("自实现与官方输出最大误差:", (a - b).abs().max().item()) # 输出 ~1e-7 量级 → 数学完全等价(仅浮点运算顺序差异)
# ---- 三个演示的调用入口(把 3.1~3.13 所有代码块按顺序拼成一个文件时会依次自动执行)----visualize_attention() # 演示 2a:注意力矩阵长什么样(行和为 1、注意力有差异)ablation() # 演示 2b:拆掉位置编码 / 残差 / 换读出,看每个部件的作用compare_with_official() # 演示 3:与官方实现逐层对照,最大误差 ~1e-7



