用PyTorch从零搭建一个Transformer模型
转载：https://blog.csdn.net/2401_85373396/article/details/153516337
Transformer总体上可分为左边的encoder部分、右边的decoder部分。目前主流大模型如deepseek、qwen、gpt等均只采用了右边的decoder部分，轻量级小模型如bert、T5等均只采用了左边的encoder部分。
左边encoder部分可分为输入部分、由N层EncoderLayer堆叠而成的编码器；右边decoder部分可分为输入部分、由N层DecoderLayer堆叠而成的解码器、输出部分。
本文中的代码引入如下，推荐安装pytorch 2.4.0 2.5.0 2.6.0，最好不要使用2.7.0，本文代码环境采用2.5.0
Python import copy import math import torch from torch import nn from torch.functional import F
输入部分：
词嵌入Embedding
词嵌入是将离散的文本符号（如单词、子词）映射为连续低维向量的过程。在PyTorch中通常通过nn.Embedding实现，它将每个词的索引转换为固定维度的稠密向量，使模型能够捕捉词语间的语义关联，为后续的注意力计算和特征提取奠定数值基础。
广义的embedding是指将自然语言字符串转换为向量的过程，各个大模型LLM均有自己的分词器和词汇表，在进入Transformer的embedding前已经将自然语言处理成了二维矩阵张量。
图示张量为[[39,1592,10,2548,5]]
1第一维（通常为batch_size）： 代表一个训练批次（batch）中包含的样本数量。为了高效训练，模型会并行处理多个句子序列。例如，一个形状为 [32, 10] 的张量，表示我们一次性处理32个独立的句子。
1第二维： 代表单个样本（即一个句子）的长度，也就是该句子中包含的单词或子词（Token）的数量。继续上面的例子，10 表示这32个句子中的每一个都被统一处理或填充为10个Token长
Plain Text
class Embedding(nn.Module):
def __init__(self, d_model: int, vocab_size: int):
"""d_model:词嵌入的维度， vocab:词表的大小"""
super(Embedding, self).__init__()
self.d_model = d_model
self.vocab_size = vocab_size
self.embedding = nn.Embedding(vocab_size, d_model)
def forward(self, x):
# 乘上权重来自于论文 3.4
return self.embedding(x) * math.sqrt(self.d_model)
位置编码器PositionalEncoding
由于Transformer无循环或卷积结构，无法直接捕捉序列的位置信息。位置编码器通过正弦和余弦函数生成与词嵌入维度相同的位置向量，将其与词嵌入向量相加，使模型能够区分不同位置的token，确保序列顺序信息在建模过程中不丢失。
位置编码根据此公式实现，由于10000乘方可能导致溢出采用exp实现，即：
Python
class PositionalEncoding(nn.Module):
def __init__(self, d_model: int, dropout: float, max_len: int = 5000):
"""d_model:词嵌入维度，dropout:置0比率，max_len:最大长度"""
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
# 实现公式 3.5, 10000^(-2i/d_model) = exp(2i × (-ln(10000)/d_model))
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):
# (batch, max_len, d_model)
# 论文不参与反向传播
x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
return self.dropout(x)
架构图中位置编码器与embedding中间有个"+"号，对应到正向传播代码就是：
x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
Encoder部分：
多头注意力MultiHeadAttention
多头注意力是Transformer的核心组件，它将输入向量通过多个并行的注意力头分别计算注意力分布，再将结果拼接融合。这种设计能让模型同时捕捉不同尺度、不同维度的语义关联（如局部依赖和全局依赖），PyTorch中可通过自定义线性层和注意力计算逻辑实现多头并行处理。
注意力机制
注意力机制的核心思想源于人类视觉和认知系统：在处理信息时，我们不会同等地对待所有输入，而是会选择性地聚焦于与当前任务最相关的部分，同时忽略其他不重要的信息。
在神经网络中，这一思想被抽象为一个可计算的数学过程。它的关键作用可以概括为 “权重分配” 。具体来说，对于一个查询（Query），模型会计算它与一系列键（Key）的相似度，然后根据这些相似度得分（即注意力权重）来对对应的值（Value）进行加权求和。这个权重就代表了模型在处理当前查询时，应该对每个值投入多少“注意力”。
Python
def attention(self, query, key, value, mask=None, dropout=None):
d_k = query.shape[-1] # d_model:词嵌入维度
# torch.matmul A的最后一个维度 必须等于B的倒数第二个维度
attn = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
attn = attn.masked_fill(mask == 0, -1e9)
# 最后一维进行softmax操作
attn = F.softmax(attn, dim=-1)
if dropout is not None:
attn = dropout(attn)
return torch.matmul(attn, value), attn
这段代码出现了概念：掩码张量（attn.masked_fill）。
1掩码张量
在 Transformer 自注意力机制中，掩码张量（Mask）是实现时序依赖约束的关键组件。由于自注意力会计算序列中所有位置间的注意力分数，为避免模型在处理当前位置时 “看到” 未来位置的信息（如语言模型生成时），掩码张量会对未来位置对应的注意力分数设置极小值（如负无穷），经过 Softmax 后这些位置的权重趋近于 0。这使得模型仅能基于当前及之前位置的信息计算注意力，保障了时序任务的合理性。
由于掩码张量计算后，还需要计算softmax。由于softmax函数中的e^x，当是负无穷的时候，值趋近于0，因此采用负无穷。
多头注意力
多头注意力由多个头来分割最后⼀维的词嵌⼊向量。
实验结果表明，Multi-head可以在更细致的层⾯上提取不同head的特征，总体计算量和单⼀head相同的情况下，提取特征的效果更佳。
Python
class MultiHeadAttention(nn.Module):
def __init__(self, head: int, d_model: int, dropout: float = 0.1):
"""head代表头数，d_model代表词嵌⼊的维度"""
super(MultiHeadAttention, self).__init__()
assert d_model % head == 0, "d_model 必须可以整除 head"
self.d_k = d_model // head# 每个头获得的分割词向量维度d_k
self.head = head
self.w_q = nn.Linear(d_model, d_model, bias=False)
self.w_k = nn.Linear(d_model, d_model, bias=False)
self.w_v = nn.Linear(d_model, d_model, bias=False)
self.w_o = nn.Linear(d_model, d_model, bias=False)
self.dropout = nn.Dropout(dropout)
# 注意力
def attention(self, query, key, value, mask=None, dropout=None):
d_k = query.shape[-1]
# torch.matmul A的最后一个维度 必须等于 B的倒数第二个维度
attn = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
attn = attn.masked_fill(mask == 0, -1e9)
# 最后一维进行softmax操作
attn = F.softmax(attn, dim=-1)
if dropout is not None:
attn = dropout(attn)
return torch.matmul(attn, value), attn
def forward(self, q, k, v, mask=None):
if mask is not None:
mask = mask.unsqueeze(0)
query = self.w_q(q)# Q
key = self.w_k(k)# K
value = self.w_v(v)# V
batch_size = query.size(0)
# 多头切割
# (batch_size, max_len, d_model) -->(batch_size, max_len, head, d_k) -->(batch_size, head, max_len, d_k)
query = query.view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
key = key.view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
value = value.view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
x, self.attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
# (batch_size, head, max_len, d_k) -> (batch_size, max_len, head * d_k) -> (batch_size, max_len, d_model)
x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)
return self.w_o(x)
最关键部分是理解清楚矩阵的结构变换：
1多头处理前：初始状态是一个三维矩阵(batch_size, max_len, d_model)，多头分割后变成了4维矩阵(batch_size, max_len, head, d_k)，为了方便多头处理将head所在维度前移(batch_size, head, max_len, d_k)
1多头处理后，又重新变回三维矩阵，与输入结构保持一致(batch_size, max_len, d_model)
前馈全连接层FeedForward
前馈全连接层对多头注意力输出的每个位置向量进行独立的非线性变换，通常采用“线性层+激活函数+线性层”的结构（如ReLU或GELU激活）。它能将注意力机制捕捉到的关联特征进一步映射到更复杂的表示空间，增强模型的非线性拟合能力。
开发实现中，一般采用带有dropout的实用公式
FFN(x) = dropout(σ(xW₁ + b₁))W₂ + b₂
其中 σ 是激活函数，在原始论文中是 ReLU
Python
class FeedForward(nn.Module):
def __init__(self, d_model: int, d_ff: int, dropout=0.1):
"""d_ff：第二个线性层的输入维度"""
super(FeedForward, self).__init__()
self.linear_1 = nn.Linear(d_model, d_ff) # 包含 W₁ 和 b₁
self.linear_2 = nn.Linear(d_ff, d_model) # 包含 W₂ 和 b₂
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# 论文中 3.3 FFN(x)
# self.linear1(x) 实际计算: x @ W₁.T + b₁
# self.linear2(...) 实际计算: ... @ W₂.T + b₂
return self.linear_2(self.dropout(F.relu(self.linear_1(x))))
规范化层LayerNorm
规范化层用于稳定模型训练过程，通过对输入张量按特征维度进行均值和方差归一化，抑制梯度消失或爆炸问题。
Python
class LayerNorm(nn.Module):
# 对每个样本不同位置的向量求均值和方差，然后进行归一化
def __init__(self, features: int, eps: float = 1e-6):
""" features, 表示词嵌⼊的维度，样本的特征数量
eps是⼀个⾜够⼩的数, 在规范化公式的分⺟中出现,防⽌分⺟为0.默认是1e-6."""
super(LayerNorm, self).__init__()
self.eps = eps
self.alpha = nn.Parameter(torch.ones(features))
self.bias = nn.Parameter(torch.zeros(features))
def forward(self, x):
# 最后一个维度的均值
mean = x.mean(dim=-1, keepdim=True)
# 最后一个维度的标准差
std = x.std(dim=-1, keepdim=True)
return self.alpha * (x - mean) / (std + self.eps) + self.bias
子层连接结构SublayerConnection
子层连接结构将每个编码器子层（如多头注意力、前馈层）的输出与输入进行残差连接，再经过层归一化处理。这种“残差+归一化”的设计能有效缓解深层模型的训练困难，让梯度更顺畅地反向传播，是Transformer能堆叠多层的关键保障。
Add & Norm模块接在每⼀个Encoder Block和Decoder Block中的每⼀个⼦层的后⾯。具体来说Add表示残差连接, Norm表示LayerNorm。
Add残差连接的作⽤：和其他神经⽹络模型中的残差连接作⽤⼀致，都是为了将信息传递的更深，增强模型的拟合能⼒。试验表明残差连接的确增强了模型的表现。
Norm的作⽤：随着⽹络层数的额增加，通过多层的计算后参数可能会出现过⼤、过⼩、⽅差变⼤等现象，这会导致学习过程出现异常，模型的收敛⾮常慢。因此对每⼀层计算后的数值进⾏规范化可以提升模型的表现。
论文中数学表达形式为: LayerNorm(x + Sublayer(x)), 其中Sublayer(x)为⼦层的输出
这里实现和论文中略有不同，先进行 norm 再进行 self-attention 或 feed-forward。也就是说，sublayer可以传多头注意力MultiHeadAttention，也只可以传前馈全连接层FeedForward。
Python
class SublayerConnection(nn.Module):
def __init__(self, features: int, dropout: float = 0.1):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(features)
self.dropout = nn.Dropout(dropout)
def forward(self, x, sublayer):
# 此处和论文中 transformer 给的图略有不同, 先进行 norm 再进行 self-attention 或 feed-forward
return x + self.dropout(sublayer(self.norm(x)))
编码器层EncoderLayer
编码器层是编码器的基本构建单元，由“多头注意力子层+子层连接”和“前馈全连接子层+子层连接”组成。每个EncoderLayer完成一次完整的特征提取与变换，通过重复堆叠多层EncoderLayer，模型能逐步捕捉序列中复杂的语义和结构信息。
Python
class EncoderLayer(nn.Module):
def __init__(self, features: int, self_attn: MultiHeadAttention,
feed_forward: FeedForward, dropout: float = 0.1):
super(EncoderLayer, self).__init__()
self.features = features
self.self_attn = self_attn
self.feed_forward = feed_forward
# 编码器层中有两个⼦层连接结构
self.sublayer = nn.ModuleList(
[SublayerConnection(features, dropout) for _ in range(2)]
)
def forward(self, x, mask):
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
return self.sublayer[1](x, self.feed_forward)
编码器Encoder
编码器由N层EncoderLayer堆叠而成。它接收经过词嵌入和位置编码处理后的源序列输入，经过多层特征提取后输出源序列的上下文表示向量，该向量包含了源序列的全局语义信息，将作为解码器的注意力输入，为目标序列生成提供依据。
Plain Text
def cloneModules(module, N):
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
Python
class Encoder(nn.Module):
def __init__(self, layer: EncoderLayer, N: int):
"""初始化函数的两个参数分别代表编码器层和编码器层的个数"""
super(Encoder, self).__init__()
self.layers = cloneModules(layer, N)
self.norm = LayerNorm(layer.features)
def forward(self, x, mask):
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
解码器层DecoderLayer
解码器的核心组件，包含三个子层：掩码自注意力层、编码器-解码器注意力层和前馈网络层，每层后都应用残差连接和规范化。
架构图中下边的多头注意力来自输入端，Q=K=V；上面的多头注意力，Q来自下面的多头注意力，K、V来自左边编码器的输出，即Q!=K=V。
Python
class DecoderLayer(nn.Module):
def __init__(self, features: int, self_attn: MultiHeadAttention, cross_attn: MultiHeadAttention,
feed_forward: FeedForward, dropout: float):
super(DecoderLayer, self).__init__()
self.features = features
self.self_attn = self_attn
self.cross_attn = cross_attn
self.feed_forward = feed_forward
self.sublayer = cloneModules(SublayerConnection(features, dropout), 3)
def forward(self, x, encoder_output, source_mask, target_mask):
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))
x = self.sublayer[1](x, lambda x: self.cross_attn(x, encoder_output, encoder_output, source_mask))
return self.sublayer[2](x, self.feed_forward)
解码器Decoder
由多个解码器层堆叠组成，逐步融合编码器输出和已生成序列信息，生成目标语言的序列表示。
Python
class Decoder(nn.Module):
def __init__(self, layer: DecoderLayer, N: int):
super(Decoder, self).__init__()
self.layers = cloneModules(layer, N)
self.norm = LayerNorm(layer.features)
def forward(self, x, encoder_output, source_mask, target_mask):
for layer in self.layers:
x = layer(x, encoder_output, source_mask, target_mask)
return self.norm(x)
解码器有2个输入，一个来自解码器的x，一个来自编码器的输出encoder_output
输出部分：
输出部分相对较简单，可合并写到一个class实现。
生成器类（线性层+softmax计算层）Generator
完成架构图中的Linear、Softmax层的计算输出。
Python
class Generator(nn.Module):
def __init__(self, d_model, vocab_size):
super(Generator, self).__init__()
self.project = nn.Linear(d_model, vocab_size)
def forward(self, x):
return F.log_softmax(self.project(x), dim=-1)
Transformer模型
整合架构图中左边encoder部分、右边encoder部分，形成整体的transformer模型。
Python
def make_transformer(source_vocab: int, target_vocab: int,
N: int = 6, d_model: int = 512, d_ff: int = 2048, head: int = 8, dropout: float = 0.1):
source_embed = Embedding(d_model, source_vocab)
target_embed = Embedding(d_model, target_vocab)
source_pos = PositionalEncoding(d_model, dropout)
target_pos = PositionalEncoding(d_model, dropout)
c = copy.deepcopy
attn = MultiHeadAttention(head, d_model)
ff = FeedForward(d_model, d_ff, dropout)
# nn.Sequential是一个顺序容器，用于将多个网络层按顺序组合在一起
model = Transformer(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
source_embed, target_embed,
source_pos, target_pos,
Generator(d_model, target_vocab))
return model
打印模型结构
编写测试代码，查看生成的模型结构
Plain Text
if __name__ == '__main__':
source_vocab = 1000
target_vocab = 1000
N = 8
transformer = make_transformer(source_vocab, target_vocab, N)
print(transformer)
# 0维：batch_size，不同的句子
# 1维：句子中的token
torch.random.manual_seed(42)
x = torch.randint(0,1000,(2,4))
mask = torch.zeros(8, 4, 4)
test_result = transformer(x, x, mask, mask)
print(test_result.shape)
打印结果如下：
SQL
Transformer(
(encoder): Encoder(
(layers): ModuleList(
(0-7): 8 x EncoderLayer(
(self_attn): MultiHeadAttention(
(w_q): Linear(in_features=512, out_features=512, bias=False)
(w_k): Linear(in_features=512, out_features=512, bias=False)
(w_v): Linear(in_features=512, out_features=512, bias=False)
(w_o): Linear(in_features=512, out_features=512, bias=False)
(dropout): Dropout(p=0.1, inplace=False)
)
(feed_forward): FeedForward(
(linear_1): Linear(in_features=512, out_features=2048, bias=True)
(linear_2): Linear(in_features=2048, out_features=512, bias=True)
(dropout): Dropout(p=0.1, inplace=False)
)
(sublayer): ModuleList(
(0-1): 2 x SublayerConnection(
(norm): LayerNorm()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
)
(norm): LayerNorm()
)
(decoder): Decoder(
(layers): ModuleList(
(0-7): 8 x DecoderLayer(
(self_attn): MultiHeadAttention(
(w_q): Linear(in_features=512, out_features=512, bias=False)
(w_k): Linear(in_features=512, out_features=512, bias=False)
(w_v): Linear(in_features=512, out_features=512, bias=False)
(w_o): Linear(in_features=512, out_features=512, bias=False)
(dropout): Dropout(p=0.1, inplace=False)
)
(cross_attn): MultiHeadAttention(
(w_q): Linear(in_features=512, out_features=512, bias=False)
(w_k): Linear(in_features=512, out_features=512, bias=False)
(w_v): Linear(in_features=512, out_features=512, bias=False)
(w_o): Linear(in_features=512, out_features=512, bias=False)
(dropout): Dropout(p=0.1, inplace=False)
)
(feed_forward): FeedForward(
(linear_1): Linear(in_features=512, out_features=2048, bias=True)
(linear_2): Linear(in_features=2048, out_features=512, bias=True)
(dropout): Dropout(p=0.1, inplace=False)
)
(sublayer): ModuleList(
(0-2): 3 x SublayerConnection(
(norm): LayerNorm()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
)
(norm): LayerNorm()
)
(source_embed): Embedding(
(embedding): Embedding(1000, 512)
)
(target_embed): Embedding(
(embedding): Embedding(1000, 512)
)
(source_pos): PositionalEncoding(
(dropout): Dropout(p=0.1, inplace=False)
)
(target_pos): PositionalEncoding(
(dropout): Dropout(p=0.1, inplace=False)
)
(generator): Generator(
(project): Linear(in_features=512, out_features=1000, bias=True)
)
)
posted on 2025-11-20 00:39 Sanny.Liu-CV&&ML 阅读(0) 评论(0) 收藏 举报
浙公网安备 33010602011771号