本文使用 Pytorch 库手动实现了传统 Transformer 模型中的多头自注意力机制、残差连接和层归一化、前馈层、编码器、解码器等子模块，进而实现了对 Transformer 模型的构建。
"""
@Title: 解析 Transformer
@Time: 2025/5/10
@Author: Michael Jie
"""
import math
import torch
import torch.nn.functional as F
from torch import nn, Tensor
# 缩放点积注意力机制 (Scaled Dot-Product Attention)
class Attention(nn.Module):
def __init__(self, causal: bool = True) -> None:
"""
注意力公式：Attention(Q, K, V) = softmax(Q · K / sqrt(d_k)) · V
Args:
causal: 是否自动生成因果掩码，默认为 True
"""
super(Attention, self).__init__()
self.causal = causal
def forward(self,
q: Tensor,
k: Tensor,
v: Tensor,
padding_mask: Tensor = None,
attn_mask: Tensor = None) -> tuple[Tensor, Tensor]:
"""
填充掩码：处理变长序列，避免填充影响注意力计算
因果掩码：防止解码器在训练时看到未来的信息
Args:
q: 查询 shape=(..., seq_len_q, d_k)
k: 键 shape=(..., seq_len_k, d_k)
v: 值 shape=(..., seq_len_k, d_v)
padding_mask: 填充掩码 shape=(..., seq_len_k)
attn_mask: 因果掩码 shape=(..., seq_len_q, seq_len_k)
Returns:
output: 输出 shape=(..., seq_len_q, d_v)
weights: 注意力权重 shape=(..., seq_len_q, seq_len_k)
"""
# 注意力分数
d_k = q.size(-1)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
# 应用填充掩码
if padding_mask is not None:
# 广播 (..., 1, seq_len_k)
scores = scores.masked_fill(padding_mask.unsqueeze(-2), float("-inf"))
# 自动生成因果掩码，优先使用自定义的因果掩码
seq_len_q, seq_len_k = q.size(-2), k.size(-2)
if self.causal and attn_mask is None:
attn_mask = torch.triu(torch.ones(seq_len_q, seq_len_k), diagonal=1).bool()
# 应用因果掩码
if attn_mask is not None:
scores = scores.masked_fill(attn_mask, float("-inf"))
# 注意力权重
weights = F.softmax(scores, dim=-1)
# 再次应用填充掩码，确保填充位置的注意力权重为 0
if padding_mask is not None:
weights = weights.masked_fill(padding_mask.unsqueeze(-2), 0)
# 乘以 v 得到输出
output = torch.matmul(weights, v)
return output, weights
# 自注意力机制 (Self Attention)
class SelfAttention(nn.Module):
def __init__(self, d_model: int = 512) -> None:
"""
自注意力机制是注意力机制的一种特殊形式，其中 Q、K、V 都来自同一输入序列，
其能够捕捉序列内部元素之间的关系，不依赖 RNN 或 CNN，直接建模长距离依赖。
Args:
d_model: 特征维度，默认为 512
"""
super(SelfAttention, self).__init__()
self.attention = Attention() # 注意力机制
# 合并 Q、K、V 的线性变换
self.linear_qkv = nn.Linear(d_model, d_model * 3)
self.linear_out = nn.Linear(d_model, d_model)
def forward(self,
x: Tensor,
padding_mask: Tensor = None,
attn_mask: Tensor = None) -> Tensor:
"""
_
Args:
x: 词嵌入 shape=(batch_size, seq_len, d_model)
padding_mask: 填充掩码 shape=(batch_size, seq_len)
attn_mask: 因果掩码 shape=(seq_len, seq_len)
Returns:
output: 输出 shape=(batch_size, seq_len, d_model)
"""
# 通过线性层同时生成 Q、K、V
qkv = self.linear_qkv(x)
q, k, v = qkv.chunk(3, dim=-1) # (batch_size, seq_len, d_model)
# 应用注意力机制
output, weights = self.attention(q, k, v, padding_mask, attn_mask)
return self.linear_out(output)
# 多头自注意力机制 (Multi-Head Self Attention)
class MultiHeadSelfAttention(nn.Module):
def __init__(self, d_model: int = 512, num_heads: int = 8) -> None:
"""
多头自注意力机制是自注意力机制的扩展，通过将输入特征分割成多个头，
每个头独立计算注意力，然后将结果拼接起来，从而提高模型的多角度表达能力。
Args:
d_model: 特征维度，默认为 512
num_heads: 头数，默认为 8
"""
super(MultiHeadSelfAttention, self).__init__()
if d_model % num_heads != 0:
raise ValueError(f"d_model must be divisible by num_heads, but got {d_model} and {num_heads}")
self.num_heads = num_heads
self.attention = Attention() # 注意力机制
# 分别对 Q、K、V 进行线性变换
self.linear_q = nn.Linear(d_model, d_model)
self.linear_k = nn.Linear(d_model, d_model)
self.linear_v = nn.Linear(d_model, d_model)
self.linear_out = nn.Linear(d_model, d_model)
def forward(self,
q: Tensor,
k: Tensor,
v: Tensor,
padding_mask: Tensor = None,
attn_mask: Tensor = None) -> Tensor:
"""
Q、K、V 在不同的自注意力模块中的来源可能不同，
在编解码器自注意力中，Q 来自解码器的输入，K、V 来自编码器的输出。
Args:
q: 查询 shape=(batch_size, seq_len, d_model)
k: 键 shape=(batch_size, seq_len / seq_len_k, d_model)
v: 值 shape=(batch_size, seq_len / seq_len_k, d_model)
padding_mask: 填充掩码 shape=(batch_size, seq_len / seq_len_k)
attn_mask: 因果掩码 shape=(seq_len / seq_len_k, seq_len / seq_len_k)
Returns:
output: 输出 shape=(batch_size, seq_len, d_model)
"""
q = self.linear_q(q)
k = self.linear_k(k)
v = self.linear_v(v)
batch_size, seq_len, seq_len_k = q.size(0), q.size(1), k.size(1)
# (batch_size, num_heads, seq_len, d_k)
q = q.view(batch_size, seq_len, self.num_heads, -1).transpose(1, 2)
k = k.view(batch_size, seq_len_k, self.num_heads, -1).transpose(1, 2)
v = v.view(batch_size, seq_len_k, self.num_heads, -1).transpose(1, 2)
# 调整掩码形状以匹配多头
if padding_mask is not None:
padding_mask = padding_mask.unsqueeze(1) # (batch_size, 1, seq_len)
if attn_mask is not None:
attn_mask = attn_mask.unsqueeze(0) # (1, seq_len, seq_len)
# 应用注意力机制
output, weights = self.attention(q, k, v, padding_mask, attn_mask)
# 拼接
output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
return self.linear_out(output)
# 残差连接和层归一化 (Add&Norm)
class AddNorm(nn.Module):
def __init__(self, d_model: int = 512) -> None:
"""
Add&Norm 层结合了两种操作：残差连接和层归一化，
可以使模型在训练过程中更加稳定，并且通过堆叠多个这样的层来构建更深的模型。
Args:
d_model: 特征维度，默认为 512
"""
super(AddNorm, self).__init__()
self.norm = nn.LayerNorm(d_model) # 层归一化
def forward(self, x: Tensor, y: Tensor) -> Tensor:
return self.norm(x + y)
# 前馈层 (FeedForward Layer)
class FeedForward(nn.Module):
def __init__(self,
input_dim: int = 512,
hidden_dim: int = 2048,
activation: str = "relu",
dropout: float = 0.1) -> None:
"""
全连接层（扩大维度） -> 激活函数 -> 全连接层（恢复原始维度）
通过非线性变换进一步提取和增强特征，使模型具备更强的模式识别和语义组合能力。
Args:
input_dim: 输入维度，默认为 512
hidden_dim: 隐藏层维度，默认为 2048
activation: 激活函数，默认为 "relu"
-支持："sigmoid", "tanh", "relu", "gelu", "leaky_relu", "elu"
dropout: 丢弃率，默认为 0.1
"""
super(FeedForward, self).__init__()
match activation: # 切换不同的激活函数
case "sigmoid":
activation = nn.Sigmoid()
case "tanh":
activation = nn.Tanh()
case "relu":
activation = nn.ReLU()
case "gelu":
activation = nn.GELU()
case "leaky_relu":
activation = nn.LeakyReLU()
case "elu":
activation = nn.ELU()
case _:
raise ValueError(f"Unsupported activation function: {activation}")
# Linear -> activation -> Dropout -> Linear
self.ffn = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
activation,
nn.Dropout(dropout),
nn.Linear(hidden_dim, input_dim),
)
def forward(self, x: Tensor) -> Tensor:
return self.ffn(x)
# 编码层
class EncoderLayer(nn.Module):
def __init__(self,
d_model: int = 512,
num_heads: int = 8,
dim_feedforward: int = 2048,
dropout: float = 0.1) -> None:
"""
MultiHeadSelfAttention -> AddNorm -> FeedForward -> AddNorm
Args:
d_model: 特征维度，默认为 512
num_heads: 头数，默认为 8
dim_feedforward: FFN 隐藏层维度，默认为 2048
dropout: 丢弃率，默认为 0.1
"""
super(EncoderLayer, self).__init__()
# 多头自注意力层
self.attn = MultiHeadSelfAttention(d_model, num_heads)
# Add&Norm 层
self.norm1 = AddNorm(d_model)
self.norm2 = AddNorm(d_model)
# 前馈层
self.ffn = FeedForward(d_model, dim_feedforward, dropout=dropout)
# 丢弃层
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
def forward(self,
x: Tensor,
padding_mask: Tensor = None,
attn_mask: Tensor = None) -> Tensor:
x = self.norm1(x, self.dropout1(self.attn(x, x, x, padding_mask, attn_mask)))
x = self.norm2(x, self.dropout2(self.ffn(x)))
return x
# 编码器
class Encoder(nn.Module):
def __init__(self, num_layers: int = 6, **params) -> None:
"""
编码器由多个编码层组成，每个编码层结构相同但并不共享参数。
Args:
num_layers: 层数，默认为 6
**params: 编码层参数，参考 EncoderLayer
"""
super(Encoder, self).__init__()
self.layers = nn.ModuleList([
EncoderLayer(**params)
for _ in range(num_layers)
])
def forward(self,
x: Tensor,
padding_mask: Tensor = None,
attn_mask: Tensor = None) -> Tensor:
for layer in self.layers: # 逐层传递
x = layer(x, padding_mask, attn_mask)
return x
# 解码层
class DecoderLayer(nn.Module):
def __init__(self,
d_model: int = 512,
num_heads: int = 8,
dim_feedforward: int = 2048,
dropout: float = 0.1) -> None:
"""
MultiHeadSelfAttention -> AddNorm -> MultiHeadSelfAttention -> AddNorm -> FeedForward -> AddNorm
Args:
d_model: 特征维度，默认为 512
num_heads: 头数，默认为 8
dim_feedforward: FFN 隐藏层维度，默认为 2048
dropout: 丢弃率，默认为 0.1
"""
super(DecoderLayer, self).__init__()
# 多头自注意力层
self.attn = MultiHeadSelfAttention(d_model, num_heads)
self.cross_attn = MultiHeadSelfAttention(d_model, num_heads)
# Add&Norm 层
self.norm1 = AddNorm(d_model)
self.norm2 = AddNorm(d_model)
self.norm3 = AddNorm(d_model)
# 前馈层
self.ffn = FeedForward(d_model, dim_feedforward, dropout=dropout)
# 丢弃层
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.dropout3 = nn.Dropout(dropout)
def forward(self,
y: Tensor,
memory: Tensor,
padding_mask_y: Tensor = None,
padding_mask_memory: Tensor = None,
attn_mask_y: Tensor = None,
attn_mask_memory: Tensor = None) -> None:
x = y
x = self.norm1(x, self.dropout1(self.attn(x, x, x, padding_mask_y, attn_mask_y)))
x = self.norm2(x, self.dropout2(self.attn(x, memory, memory, padding_mask_memory, attn_mask_memory)))
x = self.norm3(x, self.dropout3(self.ffn(x)))
return x
# 解码器
class Decoder(nn.Module):
def __init__(self, num_layers: int = 6, **params) -> None:
"""
解码器由多个解码层组成，每个解码层结构相同但并不共享参数。
Args:
num_layers: 层数，默认为 6
**params: 解码层参数，参考 DecoderLayer
"""
super(Decoder, self).__init__()
self.layers = nn.ModuleList([
DecoderLayer(**params)
for _ in range(num_layers)
])
def forward(self,
y: Tensor,
memory: Tensor,
padding_mask_y: Tensor = None,
padding_mask_memory: Tensor = None,
attn_mask_y: Tensor = None,
attn_mask_memory: Tensor = None) -> Tensor:
x = y
for layer in self.layers: # 逐层传递
x = layer(y, memory, padding_mask_y, padding_mask_memory, attn_mask_y, attn_mask_memory)
return x
# Transformer
class Transformer(nn.Module):
def __init__(self,
num_encoder_layers: int = 6,
num_decoder_layers: int = 6,
**params) -> None:
"""
transformer 是标准的编码器-解码器结构
Args:
num_encoder_layers: 编码器层数，默认为 6
num_decoder_layers: 解码器层数，默认为 6
**params: 编解码层参数，参考 EncoderLayer 和 DecoderLayer
"""
super(Transformer, self).__init__()
self.encoder = Encoder(num_encoder_layers, **params) # 编码器
self.decoder = Decoder(num_decoder_layers, **params) # 解码器
def forward(self,
x: Tensor,
y: Tensor,
padding_mask_x: Tensor = None,
padding_mask_y: Tensor = None,
padding_mask_memory: Tensor = None,
attn_mask_x: Tensor = None,
attn_mask_y: Tensor = None,
attn_mask_memory: Tensor = None) -> Tensor:
memory = self.encoder(x, padding_mask_x, attn_mask_x)
output = self.decoder(y, memory, padding_mask_y, padding_mask_memory, attn_mask_y, attn_mask_memory)
return output
if __name__ == '__main__':
# attention = Attention(True)
# t1, t2 = attention(
# torch.rand((2, 3, 64)),
# torch.rand((2, 5, 64)),
# torch.rand((2, 5, 512)),
# torch.tensor([[False, True, True, True, True],
# [False, False, False, False, True]])
# )
# print(t1.shape, t2.shape)
# self_attention = SelfAttention()
# t3 = self_attention(
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t3.shape)
# multi_head_self_attention = MultiHeadSelfAttention(num_heads=2)
# t4 = multi_head_self_attention(
# torch.rand((2, 3, 512)),
# torch.rand((2, 5, 512)),
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t4.shape)
# encoder_layer = EncoderLayer()
# t5 = encoder_layer(
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t5.shape)
# encoder = Encoder(dropout=0.2)
# t6 = encoder(
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t6.shape)
# decoder_layer = DecoderLayer()
# t7 = decoder_layer(
# torch.rand((2, 3, 512)),
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False],
# [False, False, True]]),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t7.shape)
# decoder = Decoder()
# t8 = decoder(
# torch.rand((2, 3, 512)),
# torch.rand((2, 5, 512)),
# torch.tensor([[False, False, False],
# [False, False, True]]),
# torch.tensor([[False, False, False, True, True],
# [False, False, True, True, True]])
# )
# print(t8.shape)
transformer = Transformer()
t9 = transformer(
torch.rand((2, 5, 512)),
torch.rand((2, 3, 512)),
torch.tensor([[False, False, False, True, True],
[False, False, True, True, True]]),
torch.tensor([[False, False, False],
[False, False, True]]),
)
print(t9.shape)