Transformer组件文本生成案例

一、Transformer的体系结构

下图是Transformer的体系结构：

在Transformer中，编码器和解码器的组件如下所示。

(1) 编码器

组件-1：Input Embedding

组件-2： Positional Encoding

组件-3：Multi-Head Attention

组件-4：Add & Norm

组件-5：Feed Forward

组件-6： Add Norm

(2) 解码器

组件-1：Output Embedding

组件-2： Positional Encoding

组件-3：Masked Multi-Head Attention

组件-4：Add & Norm

组件-5：Multi-Head Attention

组件-6：Add & Norm

组件-7：Feed Forward

组件-8： Add & Norm

组件-9：Linear

组件-10：softmax

二、文本生成需求
 

假设有一个简单的词汇表 { 我去图书馆看书 } ，那么当给出 "我去图书馆_"的时，预测下一个词"看书" ，我们不使用Transformer给出的API来实现，自己使用PyTorch模拟Transformer的功能来实现文本生成。

三、实现步骤

编码器实现：

# 步骤-1：Input Embedding

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# 设置随机种子确保结果可复现
torch.manual_seed(42)
np.random.seed(42)

# 步骤-1：Input Embedding
# 思想：将输入的文本转换为向量表示
# 过程：创建一个词表大小为5的嵌入层，将每个词映射为4维向量
vocab = {"我": 0, "去": 1, "图书馆": 2, "看书": 3, "<pad>": 4}
vocab_size = len(vocab)
d_model = 4  # 嵌入维度

# 创建嵌入层
input_embedding = nn.Embedding(vocab_size, d_model)

# 示例输入："我去图书馆"
input_indices = torch.tensor([vocab["我"], vocab["去"], vocab["图书馆"]])
embedded = input_embedding(input_indices)
print("步骤-1：Input Embedding 结果：")
print(embedded)

结果输出：

步骤-1：Input Embedding 结果：
tensor([[ 1.9269,  1.4873,  0.9007, -2.1055],
        [-0.7581,  1.0783,  0.8008,  1.6806],
        [ 0.3559, -0.6866, -0.4934,  0.2415]], grad_fn=<EmbeddingBackward0>)

# 步骤-2：Positional Encoding

# 思想：为序列中的每个位置添加位置信息
# 过程：使用正弦和余弦函数创建位置编码
max_seq_len = 10  # 最大序列长度

# 创建位置编码矩阵
positional_encoding = torch.zeros(max_seq_len, d_model)
position = torch.arange(0, max_seq_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

positional_encoding[:, 0::2] = torch.sin(position * div_term)
positional_encoding[:, 1::2] = torch.cos(position * div_term)

# 将位置编码添加到嵌入向量
pos = 0  #