13.Transformer——模型构建.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13.Transformer——模型构建\n",
    "\n",
    "**学习目标**\n",
    "\n",
    "1. 能用代码实现基于正弦余弦函数的位置编码\n",
    "\n",
    "2. 能用代码实现位置编码的前馈网络（FFN）\n",
    "\n",
    "3. 能用代码实现残差连接和层规范化(Add&Norm)结构\n",
    "\n",
    "4. 能用代码实现Transformer的编码器和解码器模块\n",
    "\n",
    "5. 能用代码构建Transformer模型\n",
    "\n",
    "****"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们再来回顾一下上次课程讲解的Transformer模型架构：\n",
    "\n",
    "Transformer 模型主要由编码器（Encoder）和解码器（Decoder）两部分组成，每部分都由多个相同的层（Layer）堆叠而成。标准的 Transformer 通常包括 6 层编码器和 6 层解码器。编码器负责将输入序列转化为高维表示，解码器则根据这些表示生成输出序列。\n",
    "\n",
    "<img src=\"./images/transformer.jpg\" style=\"zoom:60%;\" />\n",
    "\n",
    "（1）每个编码器都包括两个子层：多头自注意力和前馈神经网络。每个子层后都应用了残差连接和层归一化。\n",
    "\n",
    "（2）每个解码器包括以下三个子层：掩蔽多头自注意力、编码器-解码器注意力和前馈神经网络。同样，每个子层后也应用了残差连接和层归一化，以促进梯度流动和稳定训练。\n",
    "\n",
    "多头注意力（Multi-Head Attention）：Transformer 通过多头注意力机制，允许模型同时从不同的表示子空间捕捉信息，增强了模型的表达能力。\n",
    "\n",
    "位置编码（Positional Encoding）：由于 Transformer 模型本身不具备捕捉序列顺序的能力，因此引入了位置编码来提供序列中每个元素的位置信息。\n",
    "\n",
    "层归一化（Layer Normalization）和残差连接（Residual Connections）：Transformer 使用层归一化和残差连接来促进深层网络的训练，防止梯度消失或爆炸问题。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了能够顺利构建出Transformer模型，下面我们逐一构建Transformer的几个核心组件：\n",
    "\n",
    "- 位置编码（Positional Encoding）\n",
    "- 位置编码的前馈网络（Feed Forward Network）\n",
    "- 残差连接和层规范化(Add&Norm)\n",
    "- Transformer编码器（Encoder）\n",
    "- Transformer解码器（Decoder）\n",
    "\n",
    "****"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import tiktoken\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from tools import transformer\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义超参数\n",
    "vocab_size = 100256  # GPT4的词库大小\n",
    "d_model = 512          # 模型维度\n",
    "num_heads = 8          # 多头注意力头数\n",
    "num_layers = 6         # 编码器和解码器层数\n",
    "d_ffn = 2048            # 前馈神经网络隐藏层维度\n",
    "dropout = 0.1          # Dropout 概率\n",
    "max_len = 1000          # 最大序列长度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token IDs: [[42016, 78228, 80578, 0, 37046, 164, 116, 102, 84949, 222, 30590, 97565, 35287], [37046, 47551, 19000, 56386, 61056, 97565, 1811]]\n",
      "Decoded Texts: ['同志们!我踩着地雷了', '我现在开始排雷。']\n",
      "Token IDs:\n",
      " tensor([[42016, 78228, 80578,     0, 37046,   164,   116,   102, 84949,   222,\n",
      "         30590, 97565, 35287],\n",
      "        [37046, 47551, 19000, 56386, 61056, 97565,  1811,     0,     0,     0,\n",
      "             0,     0,     0]], device='cuda:0')\n",
      "词嵌入向量:\n",
      " torch.Size([2, 13, 512])\n"
     ]
    }
   ],
   "source": [
    "# 使用 tiktoken 分词器\n",
    "tiktok_encoder = tiktoken.get_encoding(\"cl100k_base\")\n",
    "torch.manual_seed(42)  # 固定随机种子\n",
    "\n",
    "texts = [\"同志们!我踩着地雷了\", \"我现在开始排雷。\"]\n",
    "\n",
    "# 将文本编码为 Token IDs\n",
    "tokens = [tiktok_encoder.encode(text) for text in texts]\n",
    "print(\"Token IDs:\", tokens)\n",
    "\n",
    "# 将 Token IDs 解码回文本\n",
    "decoded_texts = [tiktok_encoder.decode(token_ids) for token_ids in tokens]\n",
    "print(\"Decoded Texts:\", decoded_texts)\n",
    "\n",
    "# 创建嵌入层\n",
    "embedding_layer = nn.Embedding(vocab_size, d_model).to(device)\n",
    "\n",
    "# 将 Token IDs 转换为张量，并填充到相同长度\n",
    "token_ids = [torch.tensor(token_ids) for token_ids in tokens]\n",
    "token_ids = torch.nn.utils.rnn.pad_sequence(token_ids, batch_first=True).to(device)  # 填充到相同长度\n",
    "print(\"Token IDs:\\n\", token_ids)\n",
    "\n",
    "# 获取 Token Embedding\n",
    "token_embedding = embedding_layer(token_ids)  # 形状: [batch_size, seq_len, embedding_dim]\n",
    "print(\"词嵌入向量:\\n\", token_embedding.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13.1  位置编码\n",
    "\n",
    "由于Transformer模型本身不具有处理序列顺序的能力，位置编码使得模型能够理解单词在句子中的相对位置。\n",
    "\n",
    "我们可以使用不同频率的正弦和余弦函数为序列中的每个位置生成唯一的编码。位置编码的公式如下：\n",
    "\n",
    "$$\n",
    "PE_{(pos, 2i)} = \\sin\\left(\\frac{pos}{10000^{\\frac{2i}{d_\\text{model}}}}\\right)\n",
    "$$\n",
    "\n",
    "$$\n",
    "PE_{(pos, 2i+1)} = \\cos\\left(\\frac{pos}{10000^{\\frac{2i}{d_\\text{model}}}}\\right)\n",
    "$$\n",
    "\n",
    "其中：\n",
    "\n",
    "𝑃𝐸是位置编码矩阵；\n",
    "\n",
    "pos是词在序列中的绝对位置（从0开始）；\n",
    "\n",
    "𝑖是维度索引（从0开始）；\n",
    "\n",
    "$𝑑_{\\text{model}}$是模型的维度大小。\n",
    "\n",
    "对于每个维度𝑖，位置编码包含两个值：一个正弦值和一个余弦值，分别对应偶数索引和奇数索引。\n",
    "\n",
    "位置编码的目的是给模型提供每个词在序列中的相对位置信息，这样模型就可以利用这些信息来理解词与词之间的关系。位置编码通常被添加到词嵌入（Word Embeddings）中，然后一起输入到Transformer模型中。\n",
    "\n",
    "Transformer的位置编码选择三角函数的官方解释是：\n",
    "\n",
    "位置编码的每个维度都对应于一个正弦曲线。波长形成一个从2π到10000·2π的几何轨迹。我们之所以选择这个函数，是因为我们假设它可以让模型很容易地通过相对位置进行学习，因为对于任何固定的偏移量k,PEpos+k都可以表示为PEpos的线性函数。\n",
    "\n",
    "也就是说，每个维度都是波长不同的正弦波，波长范围是2π到10000·2π，选用10000这个比较大的数是因为三角函数式有周期的，在这个范围基本上，就不会出现波长一样的情况了。然后谷歌的科学家们为了让PE值的周期更长，还交替使用sin、cos来计算PE的值，就得到了最终的公式。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=1000, dropout=0.1):\n",
    "        \"\"\"\n",
    "        初始化位置编码模块。\n",
    "\n",
    "        参数:\n",
    "        - d_model: 模型的维度。\n",
    "        - max_len: 最大序列长度，默认为 1000。\n",
    "        - dropout: Dropout 的概率，默认为 0.1。\n",
    "        \"\"\"\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "        # 初始化位置编码矩阵\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0)  # 增加 batch 维度\n",
    "        self.register_buffer('pe', pe)  # 将位置编码注册为缓冲区\n",
    "\n",
    "    def forward(self, X):\n",
    "        \"\"\"\n",
    "        前向传播。\n",
    "\n",
    "        参数:\n",
    "        - X: 输入张量，形状为 (batch_size, seq_len, d_model)。\n",
    "\n",
    "        返回:\n",
    "        - 输出张量，形状为 (batch_size, seq_len, d_model)。\n",
    "        \"\"\"\n",
    "        X = X + self.pe[:, :X.size(1), :]  # 将位置编码加到输入上\n",
    "        return self.dropout(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入嵌入向量:\n",
      " tensor([[[-0.0699,  1.3706,  1.4708,  ...,  0.0000, -1.6456,  0.3474],\n",
      "         [-0.5197, -1.2748,  0.8772,  ...,  1.5548,  1.5875,  0.1868],\n",
      "         [ 0.9168, -2.0793, -0.2925,  ...,  1.2703, -0.9609,  0.6663],\n",
      "         ...,\n",
      "         [ 0.0181, -0.6386, -1.3335,  ...,  0.4053, -0.3721,  0.1877],\n",
      "         [-0.0000,  0.0208, -2.1964,  ...,  1.7799,  0.1963,  0.6811],\n",
      "         [ 0.0543, -0.0889, -0.6966,  ...,  2.0988,  0.2045,  2.0233]],\n",
      "\n",
      "        [[-0.0000,  1.9591,  1.3698,  ...,  0.5399, -0.3399,  2.6830],\n",
      "         [ 1.8715,  0.6171,  2.3728,  ...,  0.6622,  1.8515,  0.9824],\n",
      "         [ 1.1599, -0.2656,  0.0497,  ...,  0.8904,  0.0000,  0.0144],\n",
      "         ...,\n",
      "         [ 1.5365,  0.7202,  0.0000,  ...,  1.5813, -0.3754,  1.6866],\n",
      "         [ 1.0299,  1.6575, -0.0293,  ...,  1.5813, -0.3753,  1.6866],\n",
      "         [ 1.5448,  2.5902,  0.0716,  ...,  0.0000, -0.3752,  1.6866]]],\n",
      "       device='cuda:0', grad_fn=<NativeDropoutBackward0>)\n",
      "\n",
      "输入嵌入向量的形状:\n",
      " torch.Size([2, 13, 512])\n"
     ]
    }
   ],
   "source": [
    "pos_encoder = PositionalEncoding(d_model=d_model).to(device)\n",
    "# 为 Token Embedding 添加位置编码\n",
    "input_embedding = pos_encoder(token_embedding)  # 形状: [batch_size, seq_len, embedding_dim]\n",
    "print(\"输入嵌入向量:\\n\", input_embedding)\n",
    "print(\"\\n输入嵌入向量的形状:\\n\", input_embedding.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "多头自注意力形状：torch.Size([2, 13, 512])\n"
     ]
    }
   ],
   "source": [
    "# 实例化多头自注意力\n",
    "attention = transformer.MultiHeadSelfAttention(d_model, num_heads, dropout=0.5).to(device)\n",
    "attention.eval()\n",
    "# 打印多头自注意力输出的形状\n",
    "mul_att_shape = attention(input_embedding).shape\n",
    "print(f'多头自注意力形状：{mul_att_shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13.2  基于位置的前馈网络（FFN）\n",
    "\n",
    "这就是一个简单的MLP，它由两个全连接层组成。这两个全连接层的输入是来自前一层的输出和来自位置编码的向量。位置编码向量是一种映射，它将位置信息编码到输入向量中，使得模型能够捕捉到输入序列中各个位置之间的关系。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionWiseFFN(nn.Module):\n",
    "    def __init__(self, d_model, d_ffn, dropout=0.1):\n",
    "        super(PositionWiseFFN, self).__init__()\n",
    "        self.linear1 = nn.Linear(d_model, d_ffn)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.linear2 = nn.Linear(d_ffn, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.linear2(self.dropout(self.relu(self.linear1(X))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13.3  残差连接和层规范化(Add&Norm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "（1）什么是层规范化？\n",
    "\n",
    "层规范化（Layer Normalization, LayerNorm）是一种用于神经网络的归一化技术，主要用于稳定神经网络的训练过程。它与批量规范化（Batch Normalization, BatchNorm）类似，但适用的场景和实现方式有所不同。层规范化是对输入张量的**每个样本**进行归一化，而不是对整个批次进行归一化。具体来说，它对输入张量的**最后一个维度**（通常是特征维度）进行归一化，使得每个样本的特征分布更加稳定。\n",
    "\n",
    "（2）层规范化的计算公式\n",
    "\n",
    "对于输入张量x，层规范化的计算公式如下：\n",
    "\n",
    "$$\n",
    "\\text{LayerNorm}(x) = \\gamma \\left( \\frac{x - \\mu}{\\sqrt{\\sigma^2 + \\epsilon}} \\right) + \\beta\n",
    "$$\n",
    "其中：\n",
    "\n",
    "-  x ：输入张量，形状为  (N, *) ，其中  N  是样本数，*表示任意其他维度。\n",
    "\n",
    "-  $\\mu$：输入张量在最后一个维度上的均值，计算公式为：\n",
    "  $$\n",
    "  \\mu = \\frac{1}{H} \\sum_{i=1}^H x_i\n",
    "  $$\n",
    "  其中  H  是最后一个维度的大小。\n",
    "\n",
    "- $\\sigma^2$ ：输入张量在最后一个维度上的方差，计算公式为：\n",
    "  $$\n",
    "  \\sigma^2 = \\frac{1}{H} \\sum_{i=1}^H (x_i - \\mu)^2\n",
    "  $$\n",
    "  \n",
    "\n",
    "- $\\epsilon$：一个很小的常数，用于防止分母为零。\n",
    "\n",
    "- $\\gamma$  和 $\\beta$ ：可学习的缩放因子和偏移因子，形状与输入张量的最后一个维度相同。\n",
    "\n",
    "（3）层规范化与批量规范化的对比\n",
    "\n",
    "| 特性           | 批量规范化（BatchNorm）          | 层规范化（LayerNorm）                |\n",
    "| -------------- | -------------------------------- | ------------------------------------ |\n",
    "| **归一化维度** | 对批次维度进行归一化             | 对每个样本的最后一个维度进行归一化   |\n",
    "| **适用场景**   | 适用于固定大小的批次和图像数据   | 适用于变长序列和小批次数据           |\n",
    "| **计算方式**   | 依赖批次的统计信息（均值和方差） | 依赖每个样本的统计信息（均值和方差） |\n",
    "| **可学习参数** | 有可学习的缩放因子和偏移因子     | 有可学习的缩放因子和偏移因子         |\n",
    "| **性能**       | 在批次较大时效果好               | 在批次较小时效果好                   |\n",
    "\n",
    "（4）层规范化的计算示例\n",
    "\n",
    "假设输入张量  x  的形状为  (3, 4) ，表示 3 个样本，每个样本有 4 个特征。\n",
    "\n",
    "#### 输入张量\n",
    "\n",
    "$$\n",
    "x = \\begin{bmatrix}\n",
    "1 & 2 & 3 & 4 \\\\\n",
    "2 & 3 & 4 & 5 \\\\\n",
    "3 & 4 & 5 & 6 \\\\\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "\n",
    "\n",
    "#### 计算均值和方差\n",
    "\n",
    "- 对每个样本计算均值和方差：\n",
    "\n",
    "  - 样本 1：\n",
    "\n",
    "    均值 \n",
    "    $$\n",
    "    \\mu_1 = \\frac{1+2+3+4}{4} = 2.5\n",
    "    $$\n",
    "    方差 \n",
    "    $$\n",
    "    \\sigma_1^2 = \\frac{(1-2.5)^2 + (2-2.5)^2 + (3-2.5)^2 + (4-2.5)^2}{4} = 1.25\n",
    "    $$\n",
    "     \n",
    "\n",
    "  - 样本 2：\n",
    "\n",
    "    均值 \n",
    "    $$\n",
    "    \\mu_2 = \\frac{2+3+4+5}{4} = 3.5\n",
    "    $$\n",
    "    方差 \n",
    "    $$\n",
    "    \\sigma_2^2 = \\frac{(2-3.5)^2 + (3-3.5)^2 + (4-3.5)^2 + (5-3.5)^2}{4} = 1.25\n",
    "    $$\n",
    "     \n",
    "\n",
    "  - 样本 3：\n",
    "\n",
    "    均值 \n",
    "    $$\n",
    "    \\mu_3 = \\frac{3+4+5+6}{4} = 4.5\n",
    "    $$\n",
    "     方差 \n",
    "    $$\n",
    "    sigma_3^2 = \\frac{(3-4.5)^2 + (4-4.5)^2 + (5-4.5)^2 + (6-4.5)^2}{4} = 1.25\n",
    "    $$\n",
    "     \n",
    "\n",
    " 层规范化的计算结果：\n",
    "\n",
    "  $$\n",
    "  \\text{LayerNorm}(x) = \\begin{bmatrix}\n",
    "  \\frac{1-2.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{2-2.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{3-2.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{4-2.5}{\\sqrt{1.25 + \\epsilon}} \\\\\n",
    "  \\frac{2-3.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{3-3.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{4-3.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{5-3.5}{\\sqrt{1.25 + \\epsilon}} \\\\\n",
    "  \\frac{3-4.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{4-4.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{5-4.5}{\\sqrt{1.25 + \\epsilon}} & \\frac{6-4.5}{\\sqrt{1.25 + \\epsilon}} \\\\\n",
    "  \\end{bmatrix}\n",
    "  $$\n",
    "  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 残差连接和层规范化(Add&Norm)\n",
    "class AddNorm(nn.Module):\n",
    "    def __init__(self, normalized_shape, dropout):\n",
    "        \"\"\"\n",
    "        初始化 AddNorm 模块。\n",
    "\n",
    "        参数:\n",
    "        - normalized_shape: 层规范化的输入形状（通常是输入张量的最后一个维度）。\n",
    "        - dropout: Dropout 的概率（用于正则化）。\n",
    "        \"\"\"\n",
    "        super(AddNorm, self).__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.ln = nn.LayerNorm(normalized_shape)\n",
    "        \n",
    "    def forward(self, X, Y):\n",
    "        \"\"\"\n",
    "        前向传播函数。\n",
    "\n",
    "        参数:\n",
    "        - X: 输入张量（通常是上一层的输出）。\n",
    "        - Y: 需要与 X 进行残差连接的张量（通常是当前层的输出）。\n",
    "\n",
    "        返回:\n",
    "        - 经过残差连接和层规范化后的输出张量。\n",
    "        \"\"\"\n",
    "        # 对 Y 应用 Dropout，然后与 X 进行残差连接，最后进行层规范化\n",
    "        return self.ln(self.dropout(Y) + X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "该模块实现了 残差连接（Residual Connection） 和 层规范化（LayerNorm） 的功能。\n",
    "\n",
    "残差连接通过将当前层的输出与上一层的输出相加，帮助模型更好地训练深层网络。\n",
    "\n",
    "层规范化对输入进行归一化，使得模型对输入的分布更加稳定。\n",
    "\n",
    "Dropout 用于正则化，防止模型过拟合。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13.4  编码器-解码器结构"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 编码器\n",
    "\n",
    "结合上述编写的多头自注意力模块、位置编码模块、前馈网络、残差连接等模块，可以构建编码器。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_ffn, dropout):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - d_model: 模型的维度（每个词向量的维度）。\n",
    "        - num_heads: 多头注意力机制中的头数。\n",
    "        - d_ffn: 前馈神经网络中隐藏层的维度。\n",
    "        - dropout: Dropout 的概率。\n",
    "        \"\"\"\n",
    "        super(EncoderBlock, self).__init__()\n",
    "\n",
    "        # 多头自注意力层\n",
    "        self.attention = transformer.MultiHeadSelfAttention(d_model, num_heads, dropout)\n",
    "\n",
    "        # 前馈神经网络层\n",
    "        self.ffn = PositionWiseFFN(d_model, d_ffn, dropout)\n",
    "\n",
    "        # 两个 LayerNorm 层，用于归一化\n",
    "        self.norm1 = nn.LayerNorm(d_model)  # 用于自注意力层的输出\n",
    "        self.norm2 = nn.LayerNorm(d_model)  # 用于前馈神经网络的输出\n",
    "\n",
    "        # Dropout 层，用于防止过拟合\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, X, valid_lens=None):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - X: 编码器的输入，形状为 (batch_size, src_seq_len, d_model)。\n",
    "        - valid_lens: 有效序列长度，用于屏蔽填充部分，形状为 (batch_size,)。\n",
    "\n",
    "        返回:\n",
    "        - X: 编码器的输出，形状为 (batch_size, src_seq_len, d_model)。\n",
    "        \"\"\"\n",
    "        # 1. 多头自注意力机制\n",
    "        # 输入: X (源语言序列)\n",
    "        # 使用 valid_lens 屏蔽填充部分\n",
    "        attention_output = self.attention(X, valid_lens)  # (batch_size, src_seq_len, d_model)\n",
    "        # 残差连接 + Dropout\n",
    "        X = X + self.dropout(attention_output)  # (batch_size, src_seq_len, d_model)\n",
    "        # 层归一化\n",
    "        X = self.norm1(X)  # (batch_size, src_seq_len, d_model)\n",
    "\n",
    "        # 2. 前馈神经网络\n",
    "        # 输入: X (自注意力输出)\n",
    "        ffn_output = self.ffn(X)  # (batch_size, src_seq_len, d_model)\n",
    "        # 残差连接 + Dropout\n",
    "        X = X + self.dropout(ffn_output)  # (batch_size, src_seq_len, d_model)\n",
    "        # 层归一化\n",
    "        X = self.norm2(X)  # (batch_size, src_seq_len, d_model)\n",
    "\n",
    "        # 返回编码器的输出\n",
    "        return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 13, 512])\n"
     ]
    }
   ],
   "source": [
    "encoder = EncoderBlock(d_model, num_heads, d_ffn, dropout).to(device)\n",
    "\n",
    "encoder_outputs = encoder(input_embedding)\n",
    "print(encoder_outputs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. 解码器\n",
    "\n",
    "同样，我们也可以使用以上组件来构建一个transformer解码器。编码器的输出是解码器的输入。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderBlock(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_ffn, dropout):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - d_model: 模型的维度（每个词向量的维度）。\n",
    "        - num_heads: 多头注意力机制中的头数。\n",
    "        - d_ffn: 前馈神经网络中隐藏层的维度。\n",
    "        - dropout: Dropout 的概率。\n",
    "        \"\"\"\n",
    "        super(DecoderBlock, self).__init__()\n",
    "\n",
    "        # 第一个多头自注意力层（用于解码器自注意力）\n",
    "        self.attention1 = transformer.MultiHeadSelfAttention(d_model, num_heads, dropout)\n",
    "\n",
    "        # 第二个多头自注意力层（用于编码器-解码器注意力）\n",
    "        self.attention2 = transformer.MultiHeadSelfAttention(d_model, num_heads, dropout)\n",
    "\n",
    "        # 前馈神经网络层\n",
    "        self.ffn = PositionWiseFFN(d_model, d_ffn, dropout)\n",
    "\n",
    "        # 三个 LayerNorm 层，用于归一化\n",
    "        self.norm1 = nn.LayerNorm(d_model)  # 用于第一个注意力层的输出\n",
    "        self.norm2 = nn.LayerNorm(d_model)  # 用于第二个注意力层的输出\n",
    "        self.norm3 = nn.LayerNorm(d_model)  # 用于前馈神经网络的输出\n",
    "\n",
    "        # Dropout 层，用于防止过拟合\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, X, encoder_output, src_mask=None, tgt_mask=None):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - X: 解码器的输入，形状为 (batch_size, tgt_seq_len, d_model)。\n",
    "        - encoder_output: 编码器的输出，形状为 (batch_size, src_seq_len, d_model)。\n",
    "        - src_mask: 源语言的掩码，用于屏蔽填充部分，形状为 (batch_size, src_seq_len)。\n",
    "        - tgt_mask: 目标语言的掩码，用于屏蔽未来信息，形状为 (batch_size, tgt_seq_len)。\n",
    "\n",
    "        返回:\n",
    "        - X: 解码器的输出，形状为 (batch_size, tgt_seq_len, d_model)。\n",
    "        \"\"\"\n",
    "        # 1. 自注意力机制（解码器自注意力）\n",
    "        # 输入: X (目标语言序列)\n",
    "        # 使用 tgt_mask 屏蔽未来信息\n",
    "        attention_output = self.attention1(X, tgt_mask)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 残差连接 + Dropout\n",
    "        X = X + self.dropout(attention_output)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 层归一化\n",
    "        X = self.norm1(X)  # (batch_size, tgt_seq_len, d_model)\n",
    "\n",
    "        # 2. 编码器-解码器注意力机制\n",
    "        # 输入: X (解码器自注意力输出) 和 encoder_output (编码器输出)\n",
    "        # 使用 src_mask 屏蔽源语言中的填充部分\n",
    "        attention_output = self.attention2(X, src_mask)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 残差连接 + Dropout\n",
    "        X = X + self.dropout(attention_output)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 层归一化\n",
    "        X = self.norm2(X)  # (batch_size, tgt_seq_len, d_model)\n",
    "\n",
    "        # 3. 前馈神经网络\n",
    "        # 输入: X (编码器-解码器注意力输出)\n",
    "        ffn_output = self.ffn(X)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 残差连接 + Dropout\n",
    "        X = X + self.dropout(ffn_output)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 层归一化\n",
    "        X = self.norm3(X)  # (batch_size, tgt_seq_len, d_model)\n",
    "\n",
    "        # 返回解码器的输出\n",
    "        return X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 13, 512])\n"
     ]
    }
   ],
   "source": [
    "decoder = DecoderBlock(d_model, num_heads, d_ffn, dropout).to(device)\n",
    "\n",
    "outputs = decoder(encoder_outputs, input_embedding)\n",
    "print(outputs.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13.5 Transformer模型构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ffn, dropout, max_len=1000):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - src_vocab_size: 源语言词汇表的大小。\n",
    "        - tgt_vocab_size: 目标语言词汇表的大小。\n",
    "        - d_model: 模型的维度（每个词向量的维度）。\n",
    "        - num_heads: 多头注意力机制中的头数。\n",
    "        - num_layers: 编码器和解码器的层数。\n",
    "        - d_ffn: 前馈神经网络中隐藏层的维度。\n",
    "        - dropout: Dropout 的概率。\n",
    "        - max_len: 序列的最大长度，默认为 1000。\n",
    "        \"\"\"\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        # 编码器的词嵌入层，将源语言的词索引映射为 d_model 维的向量\n",
    "        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)\n",
    "\n",
    "        # 解码器的词嵌入层，将目标语言的词索引映射为 d_model 维的向量\n",
    "        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)\n",
    "\n",
    "        # 位置编码模块，为输入序列添加位置信息\n",
    "        self.positional_encoding = PositionalEncoding(d_model, max_len)\n",
    "\n",
    "        # 编码器层，由多个 EncoderBlock 组成\n",
    "        self.encoder_layers = nn.ModuleList([\n",
    "            EncoderBlock(d_model, num_heads, d_ffn, dropout) for _ in range(num_layers)\n",
    "        ])\n",
    "\n",
    "        # 解码器层，由多个 DecoderBlock 组成\n",
    "        self.decoder_layers = nn.ModuleList([\n",
    "            DecoderBlock(d_model, num_heads, d_ffn, dropout) for _ in range(num_layers)\n",
    "        ])\n",
    "\n",
    "        # 线性层，将解码器的输出映射为目标语言词汇表的大小\n",
    "        self.linear = nn.Linear(d_model, tgt_vocab_size)\n",
    "\n",
    "        # Dropout 层，用于防止过拟合\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, src, tgt, src_mask=None, tgt_mask=None):\n",
    "        \"\"\"\n",
    "        参数:\n",
    "        - src: 源语言输入序列，形状为 (batch_size, src_seq_len)。\n",
    "        - tgt: 目标语言输入序列，形状为 (batch_size, tgt_seq_len)。\n",
    "        - src_mask: 源语言的掩码，用于屏蔽填充部分，形状为 (batch_size, src_seq_len)。\n",
    "        - tgt_mask: 目标语言的掩码，用于屏蔽未来信息，形状为 (batch_size, tgt_seq_len)。\n",
    "\n",
    "        返回:\n",
    "        - output: 模型的输出，形状为 (batch_size, tgt_seq_len, tgt_vocab_size)。\n",
    "        \"\"\"\n",
    "        # 编码器部分\n",
    "        # 1. 对源语言输入进行词嵌入\n",
    "        src_embedded = self.encoder_embedding(src)  # (batch_size, src_seq_len, d_model)\n",
    "        # 2. 添加位置编码\n",
    "        src_embedded = self.positional_encoding(src_embedded)  # (batch_size, src_seq_len, d_model)\n",
    "        # 3. 应用 Dropout\n",
    "        src_embedded = self.dropout(src_embedded)  # (batch_size, src_seq_len, d_model)\n",
    "        # 4. 通过多个编码器层\n",
    "        encoder_output = src_embedded\n",
    "        for layer in self.encoder_layers:\n",
    "            encoder_output = layer(encoder_output, src_mask)  # (batch_size, src_seq_len, d_model)\n",
    "\n",
    "        # 解码器部分\n",
    "        # 1. 对目标语言输入进行词嵌入\n",
    "        tgt_embedded = self.decoder_embedding(tgt)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 2. 添加位置编码\n",
    "        tgt_embedded = self.positional_encoding(tgt_embedded)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 3. 应用 Dropout\n",
    "        tgt_embedded = self.dropout(tgt_embedded)  # (batch_size, tgt_seq_len, d_model)\n",
    "        # 4. 通过多个解码器层\n",
    "        decoder_output = tgt_embedded\n",
    "        for layer in self.decoder_layers:\n",
    "            decoder_output = layer(decoder_output, encoder_output, src_mask, tgt_mask)  # (batch_size, tgt_seq_len, d_model)\n",
    "\n",
    "        # 输出:将解码器的输出映射为目标语言词汇表的大小\n",
    "        output = self.linear(decoder_output)  # (batch_size, tgt_seq_len, tgt_vocab_size)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoded Texts: ['同志们!我踩着地雷了', '我现在开始排雷。']\n",
      "源序列:\n",
      " tensor([[42016, 78228, 80578,     0, 37046,   164,   116,   102, 84949,   222,\n",
      "         30590, 97565, 35287]], device='cuda:0')\n",
      "目标序列:\n",
      " tensor([[37046, 47551, 19000, 56386, 61056, 97565,  1811,     0,     0,     0,\n",
      "             0,     0,     0]], device='cuda:0')\n",
      "\n",
      "Transformer 输出形状:\n",
      " torch.Size([1, 13, 100256])\n"
     ]
    }
   ],
   "source": [
    "# 实例化 Transformer 模型\n",
    "transformer = Transformer(vocab_size, vocab_size, d_model, num_heads, num_layers, d_ffn, dropout, max_len).to(device)\n",
    "\n",
    "# 将 Token IDs 解码回文本\n",
    "decoded_texts = [tiktok_encoder.decode(token_ids) for token_ids in tokens]\n",
    "print(\"Decoded Texts:\", decoded_texts)\n",
    "\n",
    "src = token_ids[0].unsqueeze(0)  # (batch_size, src_seq_len)\n",
    "tgt = token_ids[1].unsqueeze(0)  # (batch_size, tgt_seq_len)\n",
    "print(\"源序列:\\n\", src)\n",
    "print(\"目标序列:\\n\", tgt)\n",
    "\n",
    "# 前向传播\n",
    "output = transformer(src, tgt)\n",
    "print(\"\\nTransformer 输出形状:\\n\", output.shape)  # (batch_size, tgt_seq_len, tgt_vocab_size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}