docs/chapter5/Transformer组件实现.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2964d025",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7293c68f",
   "metadata": {},
   "source": [
    "# 位置编码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0e49aedf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输入序列的位置编码:\n",
      "torch.Size([5, 100, 512])\n"
     ]
    }
   ],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\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).float() * (-torch.log(\n",
    "        torch.tensor(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)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.pe[:, :x.size(1)]\n",
    "        return x\n",
    "\n",
    "# 示例用法\n",
    "d_model = 512\n",
    "max_len = 100\n",
    "num_heads = 8\n",
    "\n",
    "# 位置编码\n",
    "pos_encoder = PositionalEncoding(d_model, max_len)\n",
    "\n",
    "# 示例输入序列\n",
    "input_sequence = torch.randn(5, max_len, d_model)\n",
    "\n",
    "# 应用位置编码\n",
    "input_sequence = pos_encoder(input_sequence)\n",
    "print(\"输入序列的位置编码:\")\n",
    "print(input_sequence.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55a1afd3",
   "metadata": {},
   "source": [
    "# 多头注意力"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "69f7ad4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attention_output shape: torch.Size([5, 100, 512])\n"
     ]
    }
   ],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_model, num_heads):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.d_model = d_model\n",
    "        assert d_model % num_heads == 0\n",
    "        self.depth = d_model // num_heads\n",
    "        \n",
    "        # 查询、键和值的线性投影\n",
    "        self.query_linear = nn.Linear(d_model, d_model)\n",
    "        self.key_linear = nn.Linear(d_model, d_model)\n",
    "        self.value_linear = nn.Linear(d_model, d_model)\n",
    "        \n",
    "        # 输出线性投影\n",
    "        self.output_linear = nn.Linear(d_model, d_model)\n",
    "    \n",
    "    def split_heads(self, x):\n",
    "        batch_size, seq_length, d_model = x.size()\n",
    "        return x.view(batch_size, seq_length, self.num_heads, self.depth).transpose(1, 2)\n",
    "    \n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        \n",
    "        # 线性投影\n",
    "        query = self.query_linear(query)\n",
    "        key = self.key_linear(key)\n",
    "        value = self.value_linear(value)\n",
    "        \n",
    "        # 分割头部\n",
    "        query = self.split_heads(query)\n",
    "        key = self.split_heads(key)\n",
    "        value = self.split_heads(value)\n",
    "        \n",
    "        # 缩放点积注意力\n",
    "        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.depth)\n",
    "        \n",
    "        # 如果提供了掩码，则应用掩码\n",
    "        if mask is not None:\n",
    "            scores += scores.masked_fill(mask == 0, -1e9)\n",
    "        \n",
    "        # 计算注意力权重并应用softmax\n",
    "        attention_weights = torch.softmax(scores, dim=-1)\n",
    "        \n",
    "        # 应用注意力到值\n",
    "        attention_output = torch.matmul(attention_weights, value)\n",
    "        \n",
    "        # 合并头部\n",
    "        batch_size, _, seq_length, d_k = attention_output.size()\n",
    "        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size,\n",
    "        seq_length, self.d_model)\n",
    "        \n",
    "        # 线性投影\n",
    "        attention_output = self.output_linear(attention_output)\n",
    "        \n",
    "        return attention_output\n",
    "\n",
    "# 示例用法\n",
    "d_model = 512\n",
    "max_len = 100\n",
    "num_heads = 8\n",
    "d_ff = 2048\n",
    "\n",
    "# 多头注意力\n",
    "multihead_attn = MultiHeadAttention(d_model, num_heads)\n",
    "\n",
    "# 示例输入序列\n",
    "input_sequence = torch.randn(5, max_len, d_model)\n",
    "\n",
    "# 多头注意力\n",
    "attention_output= multihead_attn(input_sequence, input_sequence, input_sequence)\n",
    "print(\"attention_output shape:\", attention_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f5275d3",
   "metadata": {},
   "source": [
    "# 前馈网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92fa36ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_sequence torch.Size([5, 100, 512])\n",
      "output_ff torch.Size([5, 100, 512])\n"
     ]
    }
   ],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff):\n",
    "        super(FeedForward, self).__init__()\n",
    "        self.linear1 = nn.Linear(d_model, d_ff)\n",
    "        self.linear2 = nn.Linear(d_ff, d_model)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 线性变换1\n",
    "        x = self.relu(self.linear1(x))\n",
    "        \n",
    "        # 线性变换2\n",
    "        x = self.linear2(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "# 示例用法\n",
    "d_model = 512\n",
    "max_len = 100\n",
    "num_heads = 8\n",
    "d_ff = 2048\n",
    "\n",
    "# 多头注意力\n",
    "multihead_attn = MultiHeadAttention(d_model, num_heads)\n",
    "\n",
    "# 前馈网络\n",
    "ff_network = FeedForward(d_model, d_ff)\n",
    "\n",
    "# 示例输入序列\n",
    "input_sequence = torch.randn(5, max_len, d_model)\n",
    "\n",
    "# 多头注意力\n",
    "attention_output= multihead_attn(input_sequence, input_sequence, input_sequence)\n",
    "\n",
    "# 前馈网络\n",
    "output_ff = ff_network(attention_output)\n",
    "print('input_sequence',input_sequence.shape)\n",
    "print(\"output_ff\", output_ff.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb1308ca",
   "metadata": {},
   "source": [
    "# 编码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "be0da170",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoder output shape: torch.Size([1, 100, 512])\n"
     ]
    }
   ],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_ff, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attention = MultiHeadAttention(d_model, num_heads)\n",
    "        self.feed_forward = FeedForward(d_model, d_ff)\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    \n",
    "    def forward(self, x, mask):\n",
    "        \n",
    "        # 自注意力层\n",
    "        attention_output= self.self_attention(x, x,\n",
    "        x, mask)\n",
    "        attention_output = self.dropout(attention_output)\n",
    "        x = x + attention_output\n",
    "        x = self.norm1(x)\n",
    "        \n",
    "        # 前馈层\n",
    "        feed_forward_output = self.feed_forward(x)\n",
    "        feed_forward_output = self.dropout(feed_forward_output)\n",
    "        x = x + feed_forward_output\n",
    "        x = self.norm2(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "d_model = 512\n",
    "max_len = 100\n",
    "num_heads = 8\n",
    "d_ff = 2048\n",
    "\n",
    "\n",
    "# 多头注意力\n",
    "encoder_layer = EncoderLayer(d_model, num_heads, d_ff, 0.1)\n",
    "\n",
    "# 示例输入序列\n",
    "input_sequence = torch.randn(1, max_len, d_model)\n",
    "\n",
    "# 多头注意力\n",
    "encoder_output= encoder_layer(input_sequence, None)\n",
    "print(\"encoder output shape:\", encoder_output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69d3b360",
   "metadata": {},
   "source": [
    "# 解码器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5f641f24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output shape: torch.Size([1, 100, 512])\n"
     ]
    }
   ],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, num_heads, d_ff, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.masked_self_attention = MultiHeadAttention(d_model, num_heads)\n",
    "        self.enc_dec_attention = MultiHeadAttention(d_model, num_heads)\n",
    "        self.feed_forward = FeedForward(d_model, d_ff)\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.norm3 = nn.LayerNorm(d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, encoder_output, src_mask, tgt_mask):\n",
    "        \n",
    "        # 掩码的自注意力层\n",
    "        self_attention_output= self.masked_self_attention(x, x, x, tgt_mask)\n",
    "        self_attention_output = self.dropout(self_attention_output)\n",
    "        x = x + self_attention_output\n",
    "        x = self.norm1(x)\n",
    "        \n",
    "        # 编码器-解码器注意力层\n",
    "        enc_dec_attention_output= self.enc_dec_attention(x, encoder_output, \n",
    "        encoder_output, src_mask)\n",
    "        enc_dec_attention_output = self.dropout(enc_dec_attention_output)\n",
    "        x = x + enc_dec_attention_output\n",
    "        x = self.norm2(x)\n",
    "        \n",
    "        # 前馈层\n",
    "        feed_forward_output = self.feed_forward(x)\n",
    "        feed_forward_output = self.dropout(feed_forward_output)\n",
    "        x = x + feed_forward_output\n",
    "        x = self.norm3(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "# 定义DecoderLayer的参数\n",
    "d_model = 512  # 模型的维度\n",
    "num_heads = 8  # 注意力头的数量\n",
    "d_ff = 2048    # 前馈网络的维度\n",
    "dropout = 0.1  # 丢弃概率\n",
    "batch_size = 1 # 批量大小\n",
    "max_len = 100  # 序列的最大长度\n",
    "\n",
    "# 定义DecoderLayer实例\n",
    "decoder_layer = DecoderLayer(d_model, num_heads, d_ff, dropout)\n",
    "\n",
    "\n",
    "src_mask = torch.rand(batch_size, max_len, max_len) > 0.5\n",
    "tgt_mask = torch.tril(torch.ones(max_len, max_len)).unsqueeze(0) == 0\n",
    "\n",
    "# 将输入张量传递到DecoderLayer\n",
    "output = decoder_layer(input_sequence, encoder_output, src_mask, tgt_mask)\n",
    "\n",
    "# 输出形状\n",
    "print(\"Output shape:\", output.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53652ca9",
   "metadata": {},
   "source": [
    "# Transformer模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0d1f4287",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 99, 5000])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff,\n",
    "    max_len, dropout):\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        # 定义编码器和解码器的词嵌入层\n",
    "        self.encoder_embedding = nn.Embedding(src_vocab_size, 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",
    "        # 定义编码器和解码器的多层堆叠\n",
    "        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout)\n",
    "        for _ in range(num_layers)])\n",
    "        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout)\n",
    "        for _ in range(num_layers)])\n",
    "\n",
    "        # 定义线性层\n",
    "        self.linear = nn.Linear(d_model, tgt_vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    # 生成掩码\n",
    "    def generate_mask(self, src, tgt):\n",
    "        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)\n",
    "        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)\n",
    "        seq_length = tgt.size(1)\n",
    "        nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length), diagonal=1)).bool()\n",
    "        tgt_mask = tgt_mask & nopeak_mask\n",
    "        return src_mask, tgt_mask\n",
    "\n",
    "    # 前向传播\n",
    "    def forward(self, src, tgt):\n",
    "        src_mask, tgt_mask = self.generate_mask(src, tgt)\n",
    "\n",
    "        # 编码器输入的词嵌入和位置编码\n",
    "        encoder_embedding = self.encoder_embedding(src)\n",
    "        en_positional_encoding = self.positional_encoding(encoder_embedding)\n",
    "        src_embedded = self.dropout(en_positional_encoding)\n",
    "\n",
    "        # 解码器输入的词嵌入和位置编码\n",
    "        decoder_embedding = self.decoder_embedding(tgt)\n",
    "        de_positional_encoding = self.positional_encoding(decoder_embedding)\n",
    "        tgt_embedded = self.dropout(de_positional_encoding)\n",
    "\n",
    "        enc_output = src_embedded\n",
    "        for enc_layer in self.encoder_layers:\n",
    "            enc_output = enc_layer(enc_output, src_mask)\n",
    "\n",
    "        dec_output = tgt_embedded\n",
    "        for dec_layer in self.decoder_layers:\n",
    "            dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)\n",
    "\n",
    "        output = self.linear(dec_output)\n",
    "        return output\n",
    "\n",
    "# 示例用法\n",
    "src_vocab_size = 5000\n",
    "tgt_vocab_size = 5000\n",
    "d_model = 512\n",
    "num_heads = 8\n",
    "num_layers = 6\n",
    "d_ff = 2048\n",
    "max_len = 100\n",
    "dropout = 0.1\n",
    "\n",
    "transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, \n",
    "d_ff, max_len, dropout)\n",
    "\n",
    "# 生成随机示例数据\n",
    "src_data = torch.randint(1, src_vocab_size, (5, max_len))  # (batch_size, seq_length)\n",
    "tgt_data = torch.randint(1, tgt_vocab_size, (5, max_len))  # (batch_size, seq_length)\n",
    "transformer(src_data, tgt_data[:, :-1]).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2bfb325",
   "metadata": {},
   "source": [
    "# Transformer 模型的训练和评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cb5d6b24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "第 1 轮：损失= 8.6675\n",
      "第 2 轮：损失= 8.1944\n",
      "第 3 轮：损失= 7.9596\n",
      "第 4 轮：损失= 7.7520\n",
      "第 5 轮：损失= 7.5324\n",
      "第 6 轮：损失= 7.2933\n",
      "第 7 轮：损失= 7.0099\n",
      "第 8 轮：损失= 6.7522\n",
      "第 9 轮：损失= 6.5018\n",
      "第 10 轮：损失= 6.2690\n",
      "\n",
      "虚拟数据的评估损失= 8.6879\n"
     ]
    }
   ],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index=0)\n",
    "optimizer = optim.Adam(transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)\n",
    "\n",
    "# 训练循环\n",
    "transformer.train()\n",
    "\n",
    "for epoch in range(10):\n",
    "    optimizer.zero_grad()\n",
    "    output = transformer(src_data, tgt_data[:, :-1])\n",
    "    loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:]\n",
    "    .contiguous().view(-1))\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    print(f\"第 {epoch+1} 轮：损失= {loss.item():.4f}\")\n",
    "\n",
    "\n",
    "# 虚拟数据\n",
    "src_data = torch.randint(1, src_vocab_size, (5, max_len))  # (batch_size, seq_length)\n",
    "tgt_data = torch.randint(1, tgt_vocab_size, (5, max_len))  # (batch_size, seq_length)\n",
    "\n",
    "# 评估循环\n",
    "transformer.eval()\n",
    "with torch.no_grad():\n",
    "    output = transformer(src_data, tgt_data[:, :-1])\n",
    "    loss = criterion(output.contiguous().view(-1, tgt_vocab_size), tgt_data[:, 1:]\n",
    "    .contiguous().view(-1))\n",
    "    print(f\"\\n虚拟数据的评估损失= {loss.item():.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0daf3c2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a6f8ab4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2fa0eff",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1ff73b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65c794a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41070a4a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}