diff --git a/example.py b/example.py index 148f479..12765f8 100644 --- a/example.py +++ b/example.py @@ -3,49 +3,63 @@ from transformer import SimpleTransformer, create_padding_mask, create_look_ahea def main(): - vocab_size = 1000 - d_model = 512 - num_heads = 8 - num_layers = 6 - d_ff = 2048 - max_seq_len = 100 + """主函数:演示Transformer模型的使用""" + # 模型超参数配置 + vocab_size = 1000 # 词汇表大小 + d_model = 512 # 模型维度(嵌入维度) + num_heads = 8 # 多头注意力的头数 + num_layers = 6 # Transformer编码器层数 + d_ff = 2048 # 前馈网络隐藏层维度 + max_seq_len = 100 # 最大序列长度 + + # 创建Transformer模型实例 model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len) - batch_size = 2 - seq_len = 10 + # 创建随机输入数据(模拟token ids) + batch_size = 2 # 批次大小 + seq_len = 10 # 序列长度 x = np.random.randint(0, vocab_size, (batch_size, seq_len)) - print("=== Simple Transformer Example ===") - print(f"Vocabulary size: {vocab_size}") - print(f"Model dimension: {d_model}") - print(f"Number of heads: {num_heads}") - print(f"Number of layers: {num_layers}") - print(f"Feed-forward dimension: {d_ff}") - print(f"Max sequence length: {max_seq_len}") + # 打印模型配置信息 + print("=== 简单Transformer示例 ===") + print(f"词汇表大小: {vocab_size}") + print(f"模型维度: {d_model}") + print(f"注意力头数: {num_heads}") + print(f"编码器层数: {num_layers}") + print(f"前馈网络维度: {d_ff}") + print(f"最大序列长度: {max_seq_len}") print() - print(f"Input shape: {x.shape}") - print(f"Input sample: {x[0]}") + # 打印输入信息 + print(f"输入形状: {x.shape}") + print(f"输入样本: {x[0]}") print() + # 执行前向传播 output = model.forward(x) - print(f"Output shape: {output.shape}") - print(f"Output sample (first 5 values): {output[0, 0, :5]}") + # 打印输出信息 + print(f"输出形状: {output.shape}") + print(f"输出样本(前5个值): {output[0, 0, :5]}") print() + # 打印模型参数统计 total_params = model.count_parameters() - print(f"Total parameters: {total_params:,}") + print(f"总参数量: {total_params:,}") print() - print("=== Attention Mask Examples ===") - padding_mask = create_padding_mask(x) - print(f"Padding mask shape: {padding_mask.shape}") + # 演示掩码创建 + print("=== 注意力掩码示例 ===") + # 填充掩码(用于忽略padding位置) + padding_mask = create_padding_mask(x) + print(f"填充掩码形状: {padding_mask.shape}") + + # 前瞻掩码(用于解码器,防止看到未来信息) look_ahead_mask = create_look_ahead_mask(seq_len) - print(f"Look-ahead mask shape: {look_ahead_mask.shape}") - print(f"Look-ahead mask sample:\n{look_ahead_mask[:5, :5]}") + print(f"前瞻掩码形状: {look_ahead_mask.shape}") + print(f"前瞻掩码示例:\n{look_ahead_mask[:5, :5]}") if __name__ == "__main__": diff --git a/transformer.py b/transformer.py index 5c3794b..a7770e8 100644 --- a/transformer.py +++ b/transformer.py @@ -2,155 +2,237 @@ 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) 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 class MultiHeadAttention: + """多头注意力机制""" def __init__(self, d_model, num_heads): - self.d_model = d_model - self.num_heads = num_heads - self.d_k = d_model // num_heads + self.d_model = d_model # 模型维度 + self.num_heads = num_heads # 注意力头数 + self.d_k = d_model // num_heads # 每个头的维度 - self.W_q = Linear(d_model, d_model) - self.W_k = Linear(d_model, d_model) - self.W_v = Linear(d_model, d_model) - self.W_o = Linear(d_model, d_model) + # 定义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) # 输出变换 def forward(self, Q, K, V, mask=None): + """ + 前向传播 + 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: 注意力掩码 + """ 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) + # 计算缩放点积注意力分数 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 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))) 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): + """ + 前向传播(残差连接 + 层归一化) + 1. 自注意力 -> 残差连接 -> 层归一化 + 2. 前馈网络 -> 残差连接 -> 层归一化 + """ + # 自注意力子层 attn_output = self.attention.forward(x, x, x, mask) x = self.norm1.forward(x + attn_output) + + # 前馈网络子层 ff_output = self.feed_forward.forward(x) x = self.norm2.forward(x + ff_output) return x class SimpleTransformer: + """简化的Transformer编码器模型""" def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len): self.d_model = d_model + # 词嵌入层 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) + pe[:, 0::2] = np.sin(position * div_term) # 偶数维度用正弦 + pe[:, 1::2] = np.cos(position * div_term) # 奇数维度用余弦 return pe def forward(self, x, mask=None): + """ + 前向传播 + 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 + 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 def create_padding_mask(seq, pad_idx=0): + """ + 创建填充掩码 + 用于忽略序列中的填充位置(padding tokens) + """ return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :] def create_look_ahead_mask(size): + """ + 创建前瞻掩码(下三角掩码) + 用于解码器中,防止位置i看到i之后的信息 + """ mask = np.triu(np.ones((size, size)), k=1).astype(bool) return (~mask).astype(float) if __name__ == "__main__": - vocab_size = 1000 - d_model = 512 - num_heads = 8 - num_layers = 6 - d_ff = 2048 - max_seq_len = 100 + # 模型超参数 + vocab_size = 1000 # 词汇表大小 + d_model = 512 # 模型维度 + num_heads = 8 # 注意力头数 + num_layers = 6 # Transformer层数 + d_ff = 2048 # 前馈网络隐藏层维度 + max_seq_len = 100 # 最大序列长度 + # 创建模型 model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len) + # 创建随机输入 batch_size = 2 seq_len = 10 x = np.random.randint(0, vocab_size, (batch_size, seq_len)) + # 前向传播 output = model.forward(x) - print(f"Input shape: {x.shape}") - print(f"Output shape: {output.shape}") - print(f"Model parameters: {model.count_parameters():,}") \ No newline at end of file + print(f"输入形状: {x.shape}") + print(f"输出形状: {output.shape}") + print(f"模型参数量: {model.count_parameters():,}") \ No newline at end of file