XingfenD/transformer_study
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

拆分transformer.py为独立模块

Browse Source 
按Transformer架构组件拆分：
- utils.py: 激活函数和掩码工具
- layers.py: 基础层（Linear, LayerNorm）
- attention.py: 多头注意力机制
- feedforward.py: 前馈网络
- block.py: Transformer编码器块
- model.py: SimpleTransformer模型
- inference.py: 推理引擎
- 删除旧的transformer.py文件
This commit is contained in:
XingfenD
2026-07-17 11:23:11 +08:00
parent c57a9bf9e6
commit 7856f5bf4c
10 changed files with 509 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files
transformer/__init__.py
+33
View File
@@ -0,0 +1,33 @@
"""
Transformer模型实现包
按架构组件拆分：
- utils: 激活函数和掩码工具
- layers: 基础层（Linear, LayerNorm）
- attention: 多头注意力机制
- feedforward: 前馈网络
- block: Transformer块
- model: SimpleTransformer模型
- inference: 推理引擎
"""
from .utils import softmax, relu, create_padding_mask, create_look_ahead_mask
from .layers import Linear, LayerNorm
from .attention import MultiHeadAttention
from .feedforward import FeedForward
from .block import TransformerBlock
from .model import SimpleTransformer
from .inference import InferenceEngine
__all__ = [
'softmax',
'relu',
'create_padding_mask',
'create_look_ahead_mask',
'Linear',
'LayerNorm',
'MultiHeadAttention',
'FeedForward',
'TransformerBlock',
'SimpleTransformer',
'InferenceEngine'
]
transformer/attention.py
+84
View File
@@ -0,0 +1,84 @@
"""
多头注意力机制模块
包含MultiHeadAttention类，支持KV cache
"""
import numpy as np
from .layers import Linear
from .utils import softmax
class MultiHeadAttention:
"""多头注意力机制（支持KV cache）"""
def __init__(self, d_model, num_heads):
self.d_model = d_model # 模型维度
self.num_heads = num_heads # 注意力头数
self.d_k = d_model // num_heads # 每个头的维度
# 定义Q、K、V的线性变换层
self.W_q = Linear(d_model, d_model) # Query变换
self.W_k = Linear(d_model, d_model) # Key变换
self.W_v = Linear(d_model, d_model) # Value变换
self.W_o = Linear(d_model, d_model) # 输出变换
# KV cache：用于推理时缓存历史Key和Value
self.k_cache = None # 缓存的Key [batch_size, num_heads, seq_len, d_k]
self.v_cache = None # 缓存的Value [batch_size, num_heads, seq_len, d_k]
def clear_cache(self):
"""清除KV cache"""
self.k_cache = None
self.v_cache = None
def forward(self, Q, K, V, mask=None, use_cache=False):
"""
前向传播
Q: Query张量 [batch_size, seq_len, d_model]
K: Key张量 [batch_size, seq_len, d_model]
V: Value张量 [batch_size, seq_len, d_model]
mask: 注意力掩码
use_cache: 是否使用KV cache（推理时设为True）
"""
batch_size = Q.shape[0]
seq_len = Q.shape[1]
# 线性变换并分割成多头
Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
# KV cache逻辑
if use_cache:
if self.k_cache is None:
# Prefill阶段：首次计算，缓存完整的K和V
self.k_cache = K
self.v_cache = V
else:
# Decode阶段：拼接历史缓存和新的K、V
self.k_cache = np.concatenate([self.k_cache, K], axis=2)
self.v_cache = np.concatenate([self.v_cache, V], axis=2)
# 使用完整的K和V进行注意力计算
K = self.k_cache
V = self.v_cache
# 计算缩放点积注意力分数
# Q: [batch, heads, q_len, d_k]
# K: [batch, heads, kv_len, d_k]（decode时kv_len > q_len）
attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k)
# 应用掩码（如果提供）
if mask is not None:
attn_scores = np.where(mask == 0, -1e9, attn_scores)
# Softmax得到注意力权重
attn_probs = softmax(attn_scores, axis=-1)
# 加权求和
attn_output = attn_probs @ V
# 拼接多头输出
attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
# 最终线性变换
output = self.W_o.forward(attn_output)
return output
transformer/block.py
+42
View File
@@ -0,0 +1,42 @@
"""
Transformer块模块
包含TransformerBlock类
"""
from .attention import MultiHeadAttention
from .feedforward import FeedForward
from .layers import LayerNorm
class TransformerBlock:
"""Transformer编码器块"""
def __init__(self, d_model, num_heads, d_ff):
# 多头自注意力层
self.attention = MultiHeadAttention(d_model, num_heads)
# 前馈网络层
self.feed_forward = FeedForward(d_model, d_ff)
# 两个层归一化
self.norm1 = LayerNorm(d_model)
self.norm2 = LayerNorm(d_model)
def forward(self, x, mask=None, use_cache=False):
"""
前向传播（残差连接 + 层归一化）
1. 自注意力 -> 残差连接 -> 层归一化
2. 前馈网络 -> 残差连接 -> 层归一化
x: 输入张量 [batch_size, seq_len, d_model]
mask: 注意力掩码
use_cache: 是否使用KV cache
"""
# 自注意力子层（Q=K=V，自注意力）
attn_output = self.attention.forward(x, x, x, mask, use_cache=use_cache)
x = self.norm1.forward(x + attn_output)
# 前馈网络子层
ff_output = self.feed_forward.forward(x)
x = self.norm2.forward(x + ff_output)
return x
def clear_cache(self):
"""清除该层的KV cache"""
self.attention.clear_cache()
transformer/feedforward.py
+20
View File
@@ -0,0 +1,20 @@
"""
前馈网络模块
包含FeedForward类
"""
from .layers import Linear
from .utils import relu
class FeedForward:
"""前馈神经网络（两层全连接网络）"""
def __init__(self, d_model, d_ff):
# 第一层：扩展维度
self.linear1 = Linear(d_model, d_ff)
# 第二层：恢复维度
self.linear2 = Linear(d_ff, d_model)
def forward(self, x):
"""前向传播：线性 -> ReLU -> 线性"""
return self.linear2.forward(relu(self.linear1.forward(x)))
transformer/inference.py
+128
View File
@@ -0,0 +1,128 @@
"""
推理引擎模块
包含InferenceEngine类，管理prefill和decode阶段
"""
import numpy as np
from .utils import softmax, create_look_ahead_mask
class InferenceEngine:
"""推理引擎：管理prefill和decode阶段"""
def __init__(self, model):
"""
初始化推理引擎
model: SimpleTransformer模型实例
"""
self.model = model
self.generated_tokens = [] # 已生成的token列表
self.current_position = 0 # 当前位置
def reset(self):
"""重置推理状态"""
self.model.clear_cache()
self.generated_tokens = []
self.current_position = 0
def prefill(self, input_ids):
"""
Prefill阶段：处理完整的输入序列
input_ids: 输入token序列 [batch_size, seq_len]
返回: 下一个token的logits [batch_size, vocab_size]
"""
self.reset()
batch_size = input_ids.shape[0]
seq_len = input_ids.shape[1]
# 创建前瞻掩码（防止看到未来信息）
mask = create_look_ahead_mask(seq_len)
# 执行prefill前向传播（使用KV cache）
output = self.model.prefill(input_ids, mask)
# 获取最后一个位置的输出（用于预测下一个token）
last_output = output[:, -1, :] # [batch_size, d_model]
# 简单的logits计算：使用嵌入矩阵的转置作为输出投影
# logits = last_output @ embedding.T
logits = last_output @ self.model.embedding.T # [batch_size, vocab_size]
# 更新状态
self.current_position = seq_len
return logits
def decode_step(self, input_token):
"""
Decode阶段：处理单个新token
input_token: 新token [batch_size, 1]
返回: 下一个token的logits [batch_size, vocab_size]
"""
# 执行decode前向传播（使用KV cache）
output = self.model.decode(input_token, self.current_position)
# 获取输出（单token，所以直接取[:, 0, :]）
last_output = output[:, 0, :] # [batch_size, d_model]
# 计算logits
logits = last_output @ self.model.embedding.T # [batch_size, vocab_size]
# 更新位置
self.current_position += 1
return logits
def generate(self, input_ids, max_new_tokens=50, temperature=1.0):
"""
自回归生成文本
input_ids: 输入token序列 [batch_size, seq_len]
max_new_tokens: 最大生成token数
temperature: 温度参数（控制随机性）
返回: 完整的生成序列 [batch_size, seq_len + max_new_tokens]
"""
batch_size = input_ids.shape[0]
# 保存原始输入
generated = input_ids.copy()
# Prefill阶段
logits = self.prefill(input_ids)
# 采样第一个生成的token
next_token = self._sample(logits, temperature)
generated = np.concatenate([generated, next_token], axis=1)
# Decode阶段：逐个生成token
for i in range(max_new_tokens - 1):
logits = self.decode_step(next_token)
next_token = self._sample(logits, temperature)
generated = np.concatenate([generated, next_token], axis=1)
# 检查是否生成了结束符（这里用0作为结束符）
if np.all(next_token == 0):
break
return generated
def _sample(self, logits, temperature=1.0):
"""
从logits中采样
logits: [batch_size, vocab_size]
temperature: 温度参数
返回: 采样的token [batch_size, 1]
"""
# 应用温度缩放
logits = logits / temperature
# 计算概率分布
probs = softmax(logits, axis=-1)
# 按概率采样
batch_size = logits.shape[0]
next_tokens = np.zeros((batch_size, 1), dtype=int)
for i in range(batch_size):
# 使用numpy的random.choice采样
next_tokens[i, 0] = np.random.choice(logits.shape[1], p=probs[i])
return next_tokens
transformer/layers.py
+36
View File
@@ -0,0 +1,36 @@
"""
基础层模块
包含Linear和LayerNorm层
"""
import numpy as np
class Linear:
"""全连接层（线性层）"""
def __init__(self, in_features, out_features):
# 使用He初始化权重
self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features)
# 偏置初始化为0
self.bias = np.zeros(out_features)
def forward(self, x):
"""前向传播：y = xW + b"""
return x @ self.weight + self.bias
class LayerNorm:
"""层归一化（Layer Normalization）"""
def __init__(self, d_model, eps=1e-6):
# 缩放参数
self.gamma = np.ones(d_model)
# 偏移参数
self.beta = np.zeros(d_model)
# 防止除零的小常数
self.eps = eps
def forward(self, x):
"""前向传播：对最后一个维度进行归一化"""
mean = np.mean(x, axis=-1, keepdims=True)
std = np.std(x, axis=-1, keepdims=True)
return self.gamma * (x - mean) / (std + self.eps) + self.beta
transformer/model.py
+117
View File
@@ -0,0 +1,117 @@
"""
Transformer模型模块
包含SimpleTransformer类
"""
import numpy as np
from .block import TransformerBlock
class SimpleTransformer:
"""简化的Transformer编码器模型（支持prefill/decode分离）"""
def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len):
self.d_model = d_model
self.vocab_size = vocab_size
# 词嵌入层
self.embedding = np.random.randn(vocab_size, d_model) * 0.01
# 位置编码
self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model)
# Transformer编码器层堆叠
self.transformer_blocks = [
TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
]
def _create_positional_encoding(self, max_seq_len, d_model):
"""
创建正弦余弦位置编码
使用不同频率的正弦和余弦函数生成位置信息
"""
pe = np.zeros((max_seq_len, d_model))
position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float)
div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model))
pe[:, 0::2] = np.sin(position * div_term) # 偶数维度用正弦
pe[:, 1::2] = np.cos(position * div_term) # 奇数维度用余弦
return pe
def clear_cache(self):
"""清除所有层的KV cache"""
for block in self.transformer_blocks:
block.clear_cache()
def prefill(self, x, mask=None):
"""
Prefill阶段：处理完整的输入序列
x: 输入序列 [batch_size, seq_len]
mask: 注意力掩码
返回: 输出张量 [batch_size, seq_len, d_model]
"""
seq_len = x.shape[1]
# 词嵌入 + 缩放
x = self.embedding[x] * np.sqrt(self.d_model)
# 加上位置编码
x = x + self.positional_encoding[:seq_len, :]
# 通过所有Transformer块，使用KV cache
for block in self.transformer_blocks:
x = block.forward(x, mask, use_cache=True)
return x
def decode(self, x, position):
"""
Decode阶段：处理单个新token
x: 新token [batch_size, 1]
position: 当前token在序列中的位置
返回: 输出张量 [batch_size, 1, d_model]
"""
# 词嵌入 + 缩放
x = self.embedding[x] * np.sqrt(self.d_model)
# 加上位置编码（使用当前位置）
x = x + self.positional_encoding[position:position+1, :]
# 通过所有Transformer块，使用KV cache
for block in self.transformer_blocks:
x = block.forward(x, use_cache=True)
return x
def forward(self, x, mask=None):
"""
前向传播（不使用KV cache，用于训练）
x: 输入序列 [batch_size, seq_len]
mask: 注意力掩码
"""
seq_len = x.shape[1]
# 词嵌入 + 缩放
x = self.embedding[x] * np.sqrt(self.d_model)
# 加上位置编码
x = x + self.positional_encoding[:seq_len, :]
# 通过所有Transformer块
for block in self.transformer_blocks:
x = block.forward(x, mask)
return x
def count_parameters(self):
"""统计模型参数数量"""
count = 0
count += self.embedding.size # 嵌入层参数
for block in self.transformer_blocks:
# 注意力层参数
count += block.attention.W_q.weight.size + block.attention.W_q.bias.size
count += block.attention.W_k.weight.size + block.attention.W_k.bias.size
count += block.attention.W_v.weight.size + block.attention.W_v.bias.size
count += block.attention.W_o.weight.size + block.attention.W_o.bias.size
# 前馈网络参数
count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size
count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size
# 层归一化参数
count += block.norm1.gamma.size + block.norm1.beta.size
count += block.norm2.gamma.size + block.norm2.beta.size
return count
transformer/utils.py
+40
View File
@@ -0,0 +1,40 @@
"""
工具函数模块
包含激活函数和掩码创建函数
"""
import numpy as np
def softmax(x, axis=-1):
"""Softmax激活函数，用于将 logits 转换为概率分布"""
# 减去最大值防止数值溢出
e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return e_x / np.sum(e_x, axis=axis, keepdims=True)
def relu(x):
"""ReLU激活函数"""
return np.maximum(0, x)
def create_padding_mask(seq, pad_idx=0):
"""
创建填充掩码
用于忽略序列中的填充位置（padding tokens）
seq: 输入序列 [batch_size, seq_len]
pad_idx: 填充token的索引
返回: 掩码 [batch_size, 1, 1, seq_len]
"""
return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :]
def create_look_ahead_mask(size):
"""
创建前瞻掩码（下三角掩码）
用于解码器中，防止位置i看到i之后的信息
size: 序列长度
返回: 掩码 [size, size]
"""
mask = np.triu(np.ones((size, size)), k=1).astype(bool)
return (~mask).astype(float)