目录
前言
碎碎念：吾愿以最不羁的才情，行最沉稳的道途。变成不假思索的身体习惯。
本文手撕transformer源码，包括几个模块：位置编码，多头注意力MHA，encoder/decoder、FFN。
为什么需要位置编码？
广播机制
注意力机制：计算注意力分数Q@KT，mask, softmax，score@v
用view拆分，用view拼接。根本没用到concat
positional encoding
python
# 定义位置编码类：继承nn.Module（PyTorch所有自定义模型的基类，必须继承）
class PositionalEncoding(nn.Module):
# 初始化函数：参数标注类型（面试代码规范，加分项）
# d_model：模型维度（如512）；max_len：序列最大长度（默认5000）；dropout：丢弃概率（防止过拟合）
def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
# 调用父类nn.Module的初始化函数：必须写，否则nn.Module的核心功能（参数管理、前向传播）失效
super().__init__()
# 定义Dropout层：训练时随机丢弃部分位置编码，防止过拟合
self.dropout = nn.Dropout(p=dropout)
# 初始化位置编码矩阵：shape=(max_len, d_model)，初始全0
# max_len行对应每个位置（0~4999），d_model列对应每个特征维度
pe = torch.zeros(max_len, d_model)
# 生成位置索引：shape=(max_len,) → 转为列向量shape=(max_len, 1)
# unsqueeze(1)：增加维度，为了和后续div_term做广播运算（面试考点：广播机制）
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 计算分母项：1/(10000^(2i/d_model))，用exp+log简化计算（避免浮点溢出，面试考点）
# torch.arange(0, d_model, 2)：取偶数索引（0,2,4...），对应正弦编码维度
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
# 偶数维度（0,2,4...）赋值正弦编码：pe[:, 0::2]是切片（起始0，步长2）
pe[:, 0::2] = torch.sin(position * div_term)
# 奇数维度（1,3,5...）赋值余弦编码：pe[:, 1::2]是切片（起始1，步长2）
# 面试考点：为什么用正弦余弦？可无限延伸，适配任意长度序列（优于可学习位置编码）
pe[:, 1::2] = torch.cos(position * div_term)
# 增加batch维度：shape=(max_len, d_model) → (1, max_len, d_model)
# 适配批量输入（多个样本），batch维度可广播
pe = pe.unsqueeze(0)
# 注册为缓冲区：pe不参与参数更新（区别于nn.Parameter），但会随模型保存/加载
# 面试考点：buffer和parameter的区别？buffer仅存储，不更新；parameter参与梯度下降
self.register_buffer('pe', pe)
# 前向传播函数：输入x是词嵌入后的张量，shape=(B, T, d_model)
# B=批量大小，T=序列长度，d_model=模型维度
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 注入位置编码：仅取前x.size(1)（当前序列长度T）的编码，避免越界
# self.pe[:, :x.size(1), :] → shape=(1, T, d_model)，和x广播相加
x = x + self.pe[:, :x.size(1), :]
# 应用Dropout后返回：训练时随机丢弃，测试时不生效
return self.dropout(x)MHA(包含attention的实现，多头在forward中实现)
python
# 定义多头注意力类：核心组件，面试必写
class MultiHeadAttention(nn.Module):
# 初始化函数：d_model=模型维度，n_head=注意力头数，dropout=丢弃概率
def __init__(self, d_model: int, n_head: int, dropout: float = 0.1):
super().__init__()
# 断言：d_model必须能被n_head整除，否则无法均匀拆分多头（面试必问！）
# 例如d_model=512，n_head=8 → 每个头维度d_k=64
assert d_model % n_head == 0, "d_model must be divisible by n_head"
# 保存核心参数，供前向传播使用
self.d_model = d_model # 模型总维度
self.n_head = n_head # 注意力头数
self.d_k = d_model // n_head # 单个头的维度
# 定义Q/K/V线性投影层：输入输出都是d_model（面试考点：为什么不直接投影到d_k？）
# 答：用1个Linear层投影到d_model，再拆分多头，比定义n_head个Linear层（每个到d_k）更简洁高效
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)
# 输出投影层：合并多头后，用Linear层融合特征（面试必问：为什么需要？）
# 答：多头拼接只是维度合并，w_o引入可训练参数，让模型学习多头特征的最优组合
self.w_o = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(p=dropout)
# 抽离缩放点积注意力为独立函数：面试时更易手写，逻辑更清晰
# 输入q/k/v：shape=(B, n_head, T, d_k)；mask：掩码矩阵（可选）
# 返回：注意力输出（shape=(B, n_head, T, d_k)）、注意力权重（shape=(B, n_head, T, T)）
def scaled_dot_product_attention(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor = None
) -> tuple[torch.Tensor, torch.Tensor]:
# 计算注意力得分：Q @ K^T / sqrt(d_k)
# q.transpose(-2,-1)：交换最后两个维度，K.shape=(B,n_head,T,d_k) → K^T=(B,n_head,d_k,T)
# 得分shape=(B, n_head, T, T)，每个元素表示第i个查询token对第j个键token的相似度
# 除以sqrt(d_k)：防止d_k过大导致得分溢出，Softmax后梯度消失（面试核心考点）
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
# 掩码处理：屏蔽不需要关注的token（填充token/未来token）
if mask is not None:
# mask.shape需匹配attn_scores：(B,1,T,T)，0位置填充-1e9（Softmax后趋近于0）
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
# 最后一维Softmax：每行和为1（面试必问：为什么dim=-1？）
# 答：dim=-1对应键token维度，每行表示单个查询token的注意力权重分配，和为1表示全部分配
attn_weights = F.softmax(attn_scores, dim=-1)
# Dropout：随机丢弃部分注意力权重，防止过拟合
attn_weights = self.dropout(attn_weights)
# 注意力权重 × V：加权求和，得到每个查询token的注意力输出
output = torch.matmul(attn_weights, v)
return output, attn_weights
# 多头注意力前向传播：输入q/k/v shape=(B, T, d_model)
def forward(
self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor = None
) -> tuple[torch.Tensor, torch.Tensor]:
# 获取批量大小B（后续拆分多头需要）
B = q.size(0)
# 步骤1：Q/K/V线性投影 → shape=(B, T, d_model)
q = self.w_q(q)
k = self.w_k(k)
v = self.w_v(v)
# 步骤2：拆分多头（面试核心维度变化！必须口述）
# view(B, -1, self.n_head, self.d_k)：(B,T,d_model) → (B,T,n_head,d_k)
# -1：自动计算序列长度T，避免硬编码
# transpose(1,2)：交换1/2维度 → (B,n_head,T,d_k)，让每个头独立计算注意力
q = q.view(B, -1, self.n_head, self.d_k).transpose(1, 2)
k = k.view(B, -1, self.n_head, self.d_k).transpose(1, 2)
v = v.view(B, -1, self.n_head, self.d_k).transpose(1, 2)
# 步骤3：调用缩放点积注意力
attn_output, attn_weights = self.scaled_dot_product_attention(q, k, v, mask)
# 步骤4：合并多头（拆分的逆操作，面试核心！）
# transpose(1,2)：(B,n_head,T,d_k) → (B,T,n_head,d_k)
# contiguous()：保证张量内存连续，否则view会报错（面试考点）
attn_output = attn_output.transpose(1, 2).contiguous()
# view(B, -1, self.d_model)：(B,T,n_head,d_k) → (B,T,d_model)，合并所有头的特征
attn_output = attn_output.view(B, -1, self.d_model)
# 步骤5：输出投影 → 融合多头特征，shape=(B,T,d_model)
output = self.w_o(attn_output)
# 返回最终输出和注意力权重（权重主要用于可视化，训练时不用）
return output, attn_weightsFFN 前馈神经网络
python
# 定义编码器层：单个编码器层包含「多头自注意力 + 前馈网络」，带残差连接和LayerNorm
class EncoderLayer(nn.Module):
# 初始化：d_model=模型维度，n_head=头数，d_ff=FFN中间维度，dropout=丢弃概率
def __init__(self, d_model: int, n_head: int, d_ff: int, dropout: float = 0.1):
super().__init__()
# 多头自注意力：编码器是自注意力（Q=K=V）
self.self_attn = MultiHeadAttention(d_model, n_head, dropout)
# 前馈网络
self.ffn = FeedForward(d_model, d_ff, dropout)
# LayerNorm：归一化（面试必问：为什么用LayerNorm？）
# 答：适配NLP（小批量/变长序列），不依赖批量维度，每个token独立归一化
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(p=dropout)
self.dropout2 = nn.Dropout(p=dropout)
# 前向传播：x=输入张量，src_mask=源序列掩码（屏蔽填充token）
def forward(self, x: torch.Tensor, src_mask: torch.Tensor = None) -> torch.Tensor:
# Pre-LN结构：先归一化，再子层，最后残差（当前主流，面试考点：Pre-LN vs Post-LN）
# Post-LN：先子层，再残差，最后归一化（原始论文结构，训练稳定性差）
# 第一步：多头自注意力 + 残差连接
# self_attn(x,x,x)：自注意力，Q=K=V=x
attn_output, _ = self.self_attn(x, x, x, src_mask)
# 残差连接：x + 注意力输出（面试必问：残差作用？）
# 答：缓解深层网络梯度消失，让梯度直接通过残差路径回传
x = x + self.dropout1(attn_output)
# LayerNorm：归一化，稳定训练
x = self.norm1(x)
# 第二步：前馈网络 + 残差连接
ffn_output = self.ffn(x)
x = x + self.dropout2(ffn_output)
x = self.norm2(x)
return xencoder：包含MHA，FFN，批量归一化
python
# 定义编码器层：单个编码器层包含「多头自注意力 + 前馈网络」，带残差连接和LayerNorm
class EncoderLayer(nn.Module):
# 初始化：d_model=模型维度，n_head=头数，d_ff=FFN中间维度，dropout=丢弃概率
def __init__(self, d_model: int, n_head: int, d_ff: int, dropout: float = 0.1):
super().__init__()
# 多头自注意力：编码器是自注意力（Q=K=V）
self.self_attn = MultiHeadAttention(d_model, n_head, dropout)
# 前馈网络
self.ffn = FeedForward(d_model, d_ff, dropout)
# LayerNorm：归一化（面试必问：为什么用LayerNorm？）
# 答：适配NLP（小批量/变长序列），不依赖批量维度，每个token独立归一化
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(p=dropout)
self.dropout2 = nn.Dropout(p=dropout)
# 前向传播：x=输入张量，src_mask=源序列掩码（屏蔽填充token）
def forward(self, x: torch.Tensor, src_mask: torch.Tensor = None) -> torch.Tensor:
# Pre-LN结构：先归一化，再子层，最后残差（当前主流，面试考点：Pre-LN vs Post-LN）
# Post-LN：先子层，再残差，最后归一化（原始论文结构，训练稳定性差）
# 第一步：多头自注意力 + 残差连接
# self_attn(x,x,x)：自注意力，Q=K=V=x
attn_output, _ = self.self_attn(x, x, x, src_mask)
# 残差连接：x + 注意力输出（面试必问：残差作用？）
# 答：缓解深层网络梯度消失，让梯度直接通过残差路径回传
x = x + self.dropout1(attn_output)
# LayerNorm：归一化，稳定训练
x = self.norm1(x)
# 第二步：前馈网络 + 残差连接
ffn_output = self.ffn(x)
x = x + self.dropout2(ffn_output)
x = self.norm2(x)
return x解码decoder（掩码自注意力 MHA，交叉注意力，前馈ffn）
python
# 定义解码器层：包含「掩码自注意力 + 交叉注意力 + 前馈网络」
class DecoderLayer(nn.Module):
def __init__(self, d_model: int, n_head: int, d_ff: int, dropout: float = 0.1):
super().__init__()
# 掩码自注意力：解码器第一层，屏蔽未来token（自回归）
self.self_attn = MultiHeadAttention(d_model, n_head, dropout)
# 交叉注意力：解码器第二层，Q=解码器输出，K/V=编码器输出（面试核心）
self.cross_attn = MultiHeadAttention(d_model, n_head, dropout)
self.ffn = FeedForward(d_model, d_ff, dropout)
# 三层LayerNorm：对应三个子层
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(p=dropout)
self.dropout2 = nn.Dropout(p=dropout)
self.dropout3 = nn.Dropout(p=dropout)
# 前向传播：
# x=解码器输入，enc_output=编码器输出，tgt_mask=目标序列掩码（自回归），src_tgt_mask=源-目标掩码
def forward(
self, x: torch.Tensor, enc_output: torch.Tensor, tgt_mask: torch.Tensor = None, src_tgt_mask: torch.Tensor = None
) -> torch.Tensor:
# 第一步：掩码自注意力（屏蔽未来token，保证自回归）
attn1_output, _ = self.self_attn(x, x, x, tgt_mask)
x = x + self.dropout1(attn1_output)
x = self.norm1(x)
# 第二步：交叉注意力（关联源序列和目标序列）
# cross_attn(x, enc_output, enc_output)：Q=x（解码器），K/V=enc_output（编码器）
attn2_output, _ = self.cross_attn(x, enc_output, enc_output, src_tgt_mask)
x = x + self.dropout2(attn2_output)
x = self.norm2(x)
# 第三步：前馈网络
ffn_output = self.ffn(x)
x = x + self.dropout3(ffn_output)
x = self.norm3(x)
return xTransformer主类的整体架构：堆叠编码器，解码器，整合词嵌入，位置编码，输出层
python
# 定义Transformer主类：堆叠编码器/解码器层，整合词嵌入、位置编码、输出层
class Transformer(nn.Module):
# 初始化参数：
# src_vocab_size=源语言词汇表大小，tgt_vocab_size=目标语言词汇表大小
# num_encoder_layers=编码器层数，num_decoder_layers=解码器层数
def __init__(
self,
src_vocab_size: int,
tgt_vocab_size: int,
d_model: int = 512,
n_head: int = 8,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
d_ff: int = 2048,
max_len: int = 5000,
dropout: float = 0.1
):
super().__init__()
self.d_model = d_model
# 词嵌入层：将词汇索引（int）转为d_model维向量
self.src_embedding = nn.Embedding(src_vocab_size, d_model)
self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
# 位置编码：共享一套（面试可选提：也可分别定义，效果差异小）
self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
# 编码器层列表：堆叠num_encoder_layers层EncoderLayer
# nn.ModuleList：可迭代的层列表（面试考点：区别于nn.Sequential？）
# 答：ModuleList仅存储层，需手动遍历调用；Sequential自动按顺序调用，灵活性低
self.encoder_layers = nn.ModuleList([
EncoderLayer(d_model, n_head, d_ff, dropout) for _ in range(num_encoder_layers)
])
# 解码器层列表：堆叠num_decoder_layers层DecoderLayer
self.decoder_layers = nn.ModuleList([
DecoderLayer(d_model, n_head, d_ff, dropout) for _ in range(num_decoder_layers)
])
# 输出线性层：将d_model维特征映射到目标词汇表大小（预测每个token的概率）
self.fc_out = nn.Linear(d_model, tgt_vocab_size)
# 生成自回归掩码：解码器专用，屏蔽未来token
# sz=目标序列长度，返回mask shape=(sz,sz)
def generate_square_subsequent_mask(self, sz: int) -> torch.Tensor:
# torch.triu：生成上三角矩阵（对角线及以上为1，以下为0）
# transpose(0,1)：转置为下三角矩阵（对角线及以下为1，以上为0）
mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
# 0位置（未来token）填充-∞，1位置（当前/过去token）填充0
# Softmax后，-∞位置权重趋近于0，无法关注未来token
mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
return mask
# Transformer前向传播：
# src=源序列（词汇索引），tgt=目标序列（词汇索引）
# src_mask=源掩码，tgt_mask=目标掩码，src_tgt_mask=源-目标掩码
def forward(
self, src: torch.Tensor, tgt: torch.Tensor, src_mask: torch.Tensor = None, tgt_mask: torch.Tensor = None, src_tgt_mask: torch.Tensor = None
) -> torch.Tensor:
# 词嵌入 + 缩放：乘以sqrt(d_model)（面试考点：为什么缩放？）
# 答：平衡词嵌入和位置编码的数值范围（位置编码范围在±1，词嵌入默认方差1）
src_emb = self.src_embedding(src) * math.sqrt(self.d_model)
tgt_emb = self.tgt_embedding(tgt) * math.sqrt(self.d_model)
# 注入位置编码：词嵌入 + 位置编码，shape=(B,T,d_model)
src_emb = self.pos_encoding(src_emb)
tgt_emb = self.pos_encoding(tgt_emb)
# 编码器前向传播：逐层调用编码器层
enc_output = src_emb
for enc_layer in self.encoder_layers:
enc_output = enc_layer(enc_output, src_mask)
# 解码器前向传播：逐层调用解码器层
dec_output = tgt_emb
for dec_layer in self.decoder_layers:
dec_output = dec_layer(dec_output, enc_output, tgt_mask, src_tgt_mask)
# 输出映射：shape=(B,T,d_model) → (B,T,tgt_vocab_size)
# 面试考点：为什么不做Softmax？
# 答：训练时用CrossEntropyLoss（内部包含Softmax），避免重复计算；推理时再做Softmax
output = self.fc_out(dec_output)
return output测试案例
python
# 主函数：仅当脚本直接运行时执行（导入时不执行）
if __name__ == "__main__":
# 1. 超参数设置（面试手写时可简化，如层数设为2）
src_vocab_size = 1000 # 源语言词汇表大小（如英语）
tgt_vocab_size = 2000 # 目标语言词汇表大小（如中文）
d_model = 512 # 模型维度（行业默认512）
n_head = 8 # 注意力头数（默认8）
num_encoder_layers = 2 # 编码器层数（简化，原始论文6层）
num_decoder_layers = 2 # 解码器层数（简化）
d_ff = 2048 # FFN中间维度（默认2048）
# 2. 实例化Transformer模型
transformer = Transformer(
src_vocab_size=src_vocab_size,
tgt_vocab_size=tgt_vocab_size,
d_model=d_model,
n_head=n_head,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
d_ff=d_ff
)
# 3. 构造测试输入（随机词汇索引）
B = 2 # 批量大小
src_seq_len = 10 # 源序列长度（如10个英语单词）
tgt_seq_len = 8 # 目标序列长度（如8个中文字符）
# src shape=(2,10)：2个样本，每个样本10个词汇索引
src = torch.randint(0, src_vocab_size, (B, src_seq_len))
# tgt shape=(2,8)：2个样本，每个样本8个词汇索引
tgt = torch.randint(0, tgt_vocab_size, (B, tgt_seq_len))
# 4. 生成解码器自回归掩码（屏蔽未来token）
tgt_mask = transformer.generate_square_subsequent_mask(tgt_seq_len)
# 5. 模型前向传播
output = transformer(src, tgt, tgt_mask=tgt_mask)
# 6. 打印维度（面试必口述维度变化！）
print(f"源输入形状: {src.shape}") # torch.Size([2, 10]) → (B, src_seq_len)
print(f"目标输入形状: {tgt.shape}") # torch.Size([2, 8]) → (B, tgt_seq_len)
print(f"模型输出形状: {output.shape}") # torch.Size([2, 8, 2000]) → (B, tgt_seq_len, tgt_vocab_size)
print("模型运行成功！")总结
- 核心组件逻辑：词嵌入 + 位置编码→编码器（自注意力 + FFN）→解码器（掩码自注意力 + 交叉注意力 + FFN）→输出线性层；
- 维度守恒 ：Transformer 所有子模块输入输出维度均为
(B,T,d_model)，保证模块可串联； - 面试必记考点：多头拆分 / 合并的维度变化、缩放点积的原因、残差连接 + LayerNorm 的作用、自回归掩码的逻辑；
- 代码规范 ：继承
nn.Module、初始化调用super().__init__()、contiguous()的使用、ModuleList的作用。
transformer八股
注入位置编码：仅取前x.size(1)（当前序列长度T）的编码，避免越界 # self.pe[:, :x.size(1), :] → shape=(1, T, d_model)，和x广播相加 x = x + self.pe[:, :x.size(1), :]
self.self_attn = MultiHeadAttention(d_model, n_head, dropout) 为什么是self.sefl.
attn_output, _ = self.self_attn(x, x, x, src_mask)
self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(p=dropout) self.dropout2 = nn.Dropout(p=dropout) 为什么要重复写。只留self.norm1和self.dropout1不行吗
残差连接：x + 注意力输出（面试必问：残差作用？）
答：缓解深层网络梯度消失，让梯度直接通过残差路径回传
LayerNorm：归一化（面试必问：为什么用LayerNorm？）