66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import numpy as np
|
|
from transformer import SimpleTransformer, create_padding_mask, create_look_ahead_mask
|
|
|
|
|
|
def main():
|
|
"""主函数：演示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)
|
|
|
|
# 创建随机输入数据（模拟token ids）
|
|
batch_size = 2 # 批次大小
|
|
seq_len = 10 # 序列长度
|
|
x = np.random.randint(0, vocab_size, (batch_size, 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"输入形状: {x.shape}")
|
|
print(f"输入样本: {x[0]}")
|
|
print()
|
|
|
|
# 执行前向传播
|
|
output = model.forward(x)
|
|
|
|
# 打印输出信息
|
|
print(f"输出形状: {output.shape}")
|
|
print(f"输出样本（前5个值）: {output[0, 0, :5]}")
|
|
print()
|
|
|
|
# 打印模型参数统计
|
|
total_params = model.count_parameters()
|
|
print(f"总参数量: {total_params:,}")
|
|
print()
|
|
|
|
# 演示掩码创建
|
|
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}")
|
|
print(f"前瞻掩码示例:\n{look_ahead_mask[:5, :5]}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |