理解Transformer架构
结合前文(链接已置顶)我们知道一个transfomer架构大致由这几部分组成:
- 使用多头注意力机制输出混合了上下文的向量。
- 将混合了上下文的向量,做残差连接(和原本的向量做加法)和LayerNorm做归一化处理。
- FNN前馈网络做深度学习。
- 对词表中的词做概率计算,输出预测。
尝试用代码实现多头注意力模块:
需要先了解,torch包链接:https://pytorch-cn.com/tutorials:PyTorch 是一个开源的深度学习框架,主要由 Meta(原 Facebook)的 AI 研究团队开发和维护,现在已经成为机器学习/深度学习领域最主流的框架之一。
下面用到的主要方法:
- torch.nn.Linear:是 PyTorch 中用来实现全连接层(线性层/仿射变换)的一个类.
- torch.matmul:两个张量的乘积。
- view方法将张量的数据重新排列,不会增加也不会减少数据。比如 12 个元素,可以变成 (3,4)、(4,3)、(2,6)、(12,) 等。
- transpose:张量的转置方法,需要注意的是转置操作只是改变了张量的步幅(stride),并没有真正移动内存中的数据,所以转置后的张量通常在内存中是不连续的。要将张量转为连续,可以调用contiguous。
import torch
import torch.nn as nn
import math
class MultiHeadAttention(nn.Module):
"""
多头注意力机制模块
"""
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除 后面拼接heads要用"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
# 定义 Q, K, V 和输出的线性变换层
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def scaled_dot_product_attention(self, Q, K, V, mask=None):
# 1. 计算注意力得分 (QK^T)
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
# 2. 应用掩码 (如果提供,主要是前文提到的Activation激活函数)
if mask is not None:
# 将掩码中为 0 的位置设置为一个非常小的负数,这样 softmax 后会接近 0
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
# 3. 计算注意力权重 (Softmax)
attn_probs = torch.softmax(attn_scores, dim=-1)
# 4. 加权求和 (权重 * V)
output = torch.matmul(attn_probs, V)
return output
# 拆分向量为多头多批次处理,提高处理效率
def split_heads(self, x):
# 将输入 x 的形状从 (batch_size, seq_length, d_model)
# 变换为 (batch_size, num_heads, seq_length, d_k)
batch_size, seq_length, d_model = x.size()
return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
def combine_heads(self, x):
# 将输入 x 的形状从 (batch_size, num_heads, seq_length, d_k)
# 变回 (batch_size, seq_length, d_model)
batch_size, num_heads, seq_length, d_k = x.size()
return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
def forward(self, Q, K, V, mask=None):
# 1. 对 Q, K, V 进行线性变换
Q = self.split_heads(self.W_q(Q))
K = self.split_heads(self.W_k(K))
V = self.split_heads(self.W_v(V))
# 2. 计算缩放点积注意力
attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
# 3. 合并多头输出并进行最终的线性变换
output = self.W_o(self.combine_heads(attn_output))
return output
测试多头函数:
# 初始化一个att的多头实例
att = MultiHeadAttention(d_model=4, num_heads=2);
# 模拟输入的Q、K、V张量
my_q = torch.rand(2, 1, 4);
my_k = torch.rand(2, 1, 4);
my_v = torch.rand(2, 1, 4);
# 调用forward函数,输出进行了最终线性变换的
theOutput = att.forward(my_q, my_k, my_v);
print('output:', theOutput, theOutput.shape);
控制台看到结果output(一个Size(2,1,4)的张量):

继续用代码实现FNN:
class PositionWiseFeedForward(nn.Module):
"""
位置前馈网络模块
"""
def __init__(self, d_model, d_ff, dropout=0.1):
super(PositionWiseFeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.dropout = nn.Dropout(dropout)
self.linear2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
def forward(self, x):
# x 形状: (batch_size, seq_len, d_model)
x = self.linear1(x)
print('larger:', x, x.shape);
x = self.relu(x)
x = self.dropout(x)
print('after activation:', x);
x = self.linear2(x)
print('recur:',x, x.shape);
# 最终输出形状: (batch_size, seq_len, d_model)
return x
测试FNN:
# d_ff通常为d_model的4倍,前文有提到
theFFN = PositionWiseFeedForward(d_model=4, d_ff=16);
# 将多头注意后的结果给FNN
theFinal = theFFN.forward(theOutput);
print("Final:", theFinal);
最终运行结果:

由结果可见,整个FNN将张量的最后一个维度从4变为16后,经过relu激活函数将负数置为0后,再将维度还原回4。
至此,按照Encoder-Decoder结构,整体代码如下:
import torch
import torch.nn as nn
import math
# --- 占位符模块,将在后续小节中实现 ---
class PositionalEncoding(nn.Module):
"""
位置编码模块,先不管pass
"""
def forward(self, x):
pass
class MultiHeadAttention(nn.Module):
"""
多头注意力机制模块,已实现
"""
def forward(self, query, key, value, mask):
pass
class PositionWiseFeedForward(nn.Module):
"""
位置前馈网络模块,已实现
"""
def forward(self, x):
pass
# --- 编码器核心层 ---
class EncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = MultiHeadAttention() # 待实现
self.feed_forward = PositionWiseFeedForward() # 待实现
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask):
# 残差连接与层归一化将在 3.1.2.4 节中详细解释
# 1. 多头自注意力
attn_output = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_output))
# 2. 前馈网络
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output))
return x
# --- 解码器核心层 ---
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(DecoderLayer, self).__init__()
self.self_attn = MultiHeadAttention() # 待实现
self.cross_attn = MultiHeadAttention() # 待实现
self.feed_forward = PositionWiseFeedForward() # 待实现
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, encoder_output, src_mask, tgt_mask):
# 1. 掩码多头自注意力 (对自己)
attn_output = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout(attn_output))
# 2. 交叉注意力 (对编码器输出)
# encoder_output作为编码器的输出传给交叉注意力头
cross_attn_output = self.cross_attn(x, encoder_output, encoder_output, src_mask)
x = self.norm2(x + self.dropout(cross_attn_output))
# 3. 前馈网络
ff_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(ff_output))
return x