当前位置：首页 > news >正文

2026年了，你还只会调用API？手把手教你从零搭建Transformer模型，硬核代码复现（含位置编码、多头注意力、残差连接全解析）

news 2026/6/9 21:58:22

大家好，我是你们的技术伙伴。👋

在2026年的今天，大语言模型（LLM）已经无处不在。然而，当我们惊叹于GPT等模型的强大时，往往容易忽略它最核心的基石——Transformer架构。

很多同学看了无数遍“Attention is All You Need”的论文，也背熟了QKV（查询、键、值）的概念，但一旦落实到代码，往往两眼发黑。

今天，我不讲那些虚无缥缈的理论，我们来一场“显微镜”级的代码实战。我们将从最基础的词嵌入开始，一行代码一行代码地构建出一个完整的Transformer模型。

准备好了吗？让我们开始吧！🚀

🧱 第一章：基石——输入部分的双重奏（词嵌入与位置编码）

Transformer抛弃了RNN的循环结构，这意味着它无法天然地感知序列的顺序。为了解决这个问题，输入部分的设计至关重要，它由词嵌入（Word Embedding）和位置编码（Positional Encoding）两部分组成。

1. 词嵌入层 (Word Embedding)

这是将离散的单词转化为连续向量的第一步。我们在代码中不仅实现了嵌入，还加入了一个关键的缩放技巧。

import torch import torch.nn as nn import math class Embeddings(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.vocab_size = vocab_size self.d_model = d_model # 核心嵌入层 self.embed = nn.Embedding(vocab_size, d_model) def forward(self, x): # 关键点：乘以 sqrt(d_model) # 目的：为了平衡梯度，避免因维度变化导致的数值过大或过小，防止梯度消失或爆炸。 return self.embed(x) * math.sqrt(self.d_model)

💡 原理解析：
为什么要乘以sqrt(d_model)？这是为了在初始阶段保持向量的方差稳定，确保模型训练的平滑性。

2. 位置编码 (Positional Encoding)

这是Transformer的灵魂所在。我们使用正弦和余弦函数生成位置编码，让模型能够感知单词在序列中的位置。

class PositionEncoding(nn.Module): def __init__(self, d_model, dropout, max_len=60): super().__init__() self.dropout = nn.Dropout(p=dropout) # 1. 初始化位置编码矩阵 [max_len, d_model] pe = torch.zeros(max_len, d_model) # 2. 生成位置序列 [0, max_len) position = torch.arange(0, max_len).unsqueeze(1) # [max_len, 1] # 3. 核心公式：计算分母中的指数项 # div_term = 1 / (10000^(2i/d_model)) div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) # 4. 偶数位置用sin, 奇数位置用cos pe[:, 0::2] = torch.sin(position * div_term) # 偶数维度 pe[:, 1::2] = torch.cos(position * div_term) # 奇数维度 # 5. 增加批次维度并注册为缓冲区(不更新参数) pe = pe.unsqueeze(0) # [1, max_len, d_model] self.register_buffer('pe', pe) def forward(self, x): # 将词嵌入与位置编码相加 x = x + self.pe[:, :x.size(1)] return self.dropout(x)

可视化：
运行代码中的绘图函数，你会看到不同维度的位置编码呈现出不同的波形，这保证了每个位置都有独一无二的“指纹”。

🧠 第二章：核心——多头注意力机制 (Multi-Head Attention)

Attention机制是Transformer的“眼睛”。它让模型在处理某个词时，能够关注到句子中其他相关的词。

1. 缩放点积注意力 (Scaled Dot-Product Attention)

这是最底层的计算单元。注意其中的Mask机制，它用于遮挡未来的信息（在解码器中）或填充的无效信息。

def attention(query, key, value, mask=None, dropout=None): # 1. 计算Q和K的点积，并缩放 d_k = query.size()[-1] scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) # 2. Mask机制：将mask位置填充为极小值(-1e9)，确保softmax后概率为0 if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # 3. 计算注意力权重 p_attn = F.softmax(scores, dim=-1) # 4. 加权求和得到输出 if dropout is not None: p_attn = dropout(p_attn) return torch.matmul(p_attn, value), p_attn

2. 多头注意力 (Multi-Head Attention)

单个注意力头可能只关注局部信息，多个头可以并行关注不同子空间的信息。

class MultiHeadAttention(nn.Module): def __init__(self, embed_dim, head, dropout_p=0.1): super().__init__() assert embed_dim % head == 0 # 确保能整除 self.d_k = embed_dim // head self.head = head # 4个线性层：分别用于 Q, K, V 的投影 和 最终输出的投影 self.linears = clones(nn.Linear(embed_dim, embed_dim), 4) self.dropout = nn.Dropout(dropout_p) self.atten = None def forward(self, query, key, value, mask=None): if mask is not None: mask = mask.unsqueeze(0) # 增加头维度 batch = query.size(0) # 1. 线性变换 + 分头 # 将 [Batch, Seq_Len, Embed_Dim] -> [Batch, Head, Seq_Len, d_k] query, key, value = [ model(x).view(batch, -1, self.head, self.d_k).transpose(1, 2) for model, x in zip(self.linears, (query, key, value)) ] # 2. 调用底层attention函数 x, self.atten = attention(query, key, value, mask, self.dropout) # 3. 合并多头结果 # [Batch, Head, Seq_Len, d_k] -> [Batch, Seq_Len, Embed_Dim] x = x.transpose(1, 2).contiguous().view(batch, -1, self.head * self.d_k) # 4. 最终线性投影 return self.linears[-1](x)

代码解析：
这里的clones函数使用了nn.ModuleList深拷贝多个线性层，这是为了确保参数不共享。

⚙️ 第三章：骨架——编码器与解码器

有了底层积木，现在我们来搭建模型的主体。

1. 编码器层 (Encoder Layer)

编码器层由两个子层组成：多头自注意力机制和前馈全连接网络（FFN）。每个子层都使用了残差连接（Residual Connection）和层规范化（Layer Normalization）。

class EncoderLayer(nn.Module): def __init__(self, d_model, self_attn, feed_forward, dropout=0.1): super().__init__() self.d_model = d_model self.self_attn = self_attn self.feed_forward = feed_forward # 克隆两个子层连接结构（包含残差连接和规范化） self.sublayer = clones(SublayerConnection(d_model, dropout), 2) def forward(self, x, mask): # 第1个子层：多头自注意力 x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第2个子层：前馈全连接网络 x = self.sublayer[1](x, lambda x: self.feed_forward(x)) return x

2. 解码器层 (Decoder Layer)

解码器比编码器更复杂，它多了一个编码器-解码器注意力层，用于关注源序列的信息。

class DecoderLayer(nn.Module): def __init__(self, d_model, self_attn, src_attn, feed_forward, dropout=0.1): super().__init__() self.d_model = d_model self.self_attn = self_attn # 掩码多头注意力（防止看到未来信息） self.src_attn = src_attn # 编码器-解码器注意力 self.feed_forward = feed_forward # 3个子层连接 self.layers = clones(SublayerConnection(d_model, dropout), mit 3) def forward(self, x, encoder_output, source_mask, target_mask): # 第1子层：掩码自注意力（Q=K=V=x） x = self.layers[0](x, lambda x: self.self_attn(x, x, x, target_mask)) # 第2子层：编码器-解码器注意力 # Q来自解码器，K和V来自编码器输出 x = self.layers[1](x, lambda x: self.src_attn(x, encoder_output, encoder_output, source_mask)) # 第3子层：前馈网络 x = self.layers[2](x, self.feed_forward) return x

🏁 第四章：输出——生成最终结果

模型的最后一步是将解码器输出的特征向量转化为具体的词汇概率分布。

1. 生成器 (Generator)

使用一个线性层将特征映射到词汇表大小，然后通过LogSoftmax转化为对数概率。

class Generator(nn.Module): def __init__(self, d_model, vocab_size): super().__init__() # 线性层：将解码器输出的特征转为词汇表大小的分数 self.linear = nn.Linear(d_model, vocab_size) def forward(self, x): # LogSoftmax + NLLLoss 是分类任务的经典组合 return F.log_softmax(self.linear(x), dim=-1)

2. 组装完整的Transformer

最后，我们将所有组件拼装起来，形成一个完整的EncoderDecoder类。

class EncoderDecoder(nn.Module): def __init__(self, source_embed, encoder, target_embed, decoder, generator): super().__init__() self.source_embed = source_embed self.encoder = encoder self.target_embed = target_embed self.decoder = decoder self.generator = generator def forward(self, source_x, target_y, source_mask, target_mask): # 1. 编码 encoder_result = self.encode(source_x, source_mask) # 2. 解码 decoder_result = self.decode(target_y, encoder_result, source_mask, target_mask) # 3. 生成最终输出 return self.generator(decoder_result) def encode(self, source_x, source_mask): embed_x = self.source_embed(source_x) return self.encoder(embed_x, source_mask) def decode(self, target_y, encoder_output, source_mask, target_mask): embed_y = self.target_embed(target_y) return self.decoder(embed_y, encoder_output, source_mask, target_mask)