transformer_torch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c73fdb87-4574-4dd4-a8ae-67b4040fa3e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 后续代码所需的依赖\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math, copy, time\n",
    "\n",
    "class TokenEmbedding(nn.Embedding):\n",
    "    \"\"\"\n",
    "    使用torch.nn的Embedding模块\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, d_model):\n",
    "        \"\"\"\n",
    "        TokenEmbedding类\n",
    "\n",
    "        :param vocab_size: 词汇表的大小\n",
    "        :param d_model: 模型的维度\n",
    "        :padding的索引为1，即token索引为1时，Embedding补0\n",
    "        \"\"\"\n",
    "        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "956b9cb9-3aab-481b-88dd-34b92f9563d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模块中的参数数量为: 512000\n"
     ]
    }
   ],
   "source": [
    "tok_emb = TokenEmbedding(1000, 512)\n",
    "num_params = sum(p.numel() for p in tok_emb.parameters())\n",
    "print(\"模块中的参数数量为:\", num_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf5a9738-b4d0-4ab1-be42-108c6fffb969",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "res: tensor([[[ 0.1994,  0.9800,  1.3162,  ...,  0.5627,  0.4720,  0.9675],\n",
      "         [ 0.5742, -0.6087,  0.5124,  ...,  1.2031,  0.7531, -0.6712],\n",
      "         [ 1.5465, -0.5008,  0.9264,  ...,  0.0815, -0.2846,  0.0975]],\n",
      "\n",
      "        [[ 0.7377, -0.0586, -1.3138,  ...,  0.4524, -0.2046,  1.6616],\n",
      "         [ 1.1249, -1.7420, -1.6807,  ..., -0.0495,  0.3476,  1.1462],\n",
      "         [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000]]],\n",
      "       grad_fn=<EmbeddingBackward0>)\n",
      "res.shape: torch.Size([2, 3, 512])\n"
     ]
    }
   ],
   "source": [
    "# x是batch_size为2, seq_len为3，索引为1的会被padding为0\n",
    "x = torch.LongTensor([[6, 5, 4], [3, 2, 1]])\n",
    "res = tok_emb(x)\n",
    "print(\"res:\", res)\n",
    "print(\"res.shape:\", res.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a45fb17-e01b-434e-9003-f8193d8beb6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    \"\"\"\n",
    "    计算正余弦位置编码。\n",
    "    \"\"\"\n",
    "    def __init__(self, d_model, max_len):\n",
    "        \"\"\"\n",
    "        正余弦位置编码类\n",
    "\n",
    "        :param d_model: 模型的维度\n",
    "        :param max_len: 最大序列长度\n",
    "        \"\"\"\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "\n",
    "        # 初始化位置编码矩阵\n",
    "        self.encoding = torch.zeros(max_len, d_model)\n",
    "        self.encoding.requires_grad = False  # 不需要计算梯度\n",
    "\n",
    "        pos = torch.arange(0, max_len)\n",
    "        pos = pos.float().unsqueeze(dim=1)\n",
    "\n",
    "        # 'i'表示d_model的索引（例如，嵌入大小=50，'i' = [0,50]）\n",
    "        # “step=2”表示将'i'乘以二（与2 * i相同）\n",
    "        _2i = torch.arange(0, d_model, step=2).float()\n",
    "        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))\n",
    "        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # self.encoding\n",
    "        # [max_len = 512, d_model = 512]\n",
    "\n",
    "        batch_size, seq_len = x.size()\n",
    "        # [batch_size = 8, seq_len = 30]\n",
    "\n",
    "        return self.encoding[:seq_len, :]\n",
    "        # [seq_len = 30, d_model = 512]\n",
    "        # 将与 tok_emb 相加：[8, 30, 512]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2f2e1eb5-18eb-4b50-92c3-0ed6c5d29f01",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模块中的参数数量为: 0\n"
     ]
    }
   ],
   "source": [
    "pe = PositionalEncoding(512,512)\n",
    "num_params = sum(p.numel() for p in pe.parameters())\n",
    "print(\"模块中的参数数量为:\", num_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6d60b122-114f-48cf-b241-fd57766e0afe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "res: tensor([[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
      "          0.0000e+00,  1.0000e+00],\n",
      "        [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,\n",
      "          1.0366e-04,  1.0000e+00],\n",
      "        [ 9.0930e-01, -4.1615e-01,  9.3641e-01,  ...,  1.0000e+00,\n",
      "          2.0733e-04,  1.0000e+00]])\n",
      "res.shape: torch.Size([3, 512])\n"
     ]
    }
   ],
   "source": [
    "# x是batch_size为2, seq_len为3\n",
    "x = torch.LongTensor([[6, 5, 4], [3, 2, 1]])\n",
    "res = pe.forward(x)\n",
    "print(\"res:\", res)\n",
    "# 返回的形状是[seq_len = 3, d_model = 512]\n",
    "print(\"res.shape:\", res.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5cf833ab-9ab5-47ec-923e-3e812c7adfeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    token embedding + positional encoding\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, d_model, max_len, drop_prob):\n",
    "        \"\"\"\n",
    "        包含Embedding和位置编码的类\n",
    "\n",
    "        :param vocab_size: 词汇表大小\n",
    "        :param d_model: 模型的维度\n",
    "        :param max_len: 最大序列长度\n",
    "        :param drop_prob: dropout 正则化概率，防止过拟合\n",
    "        \"\"\"\n",
    "        super(TransformerEmbedding, self).__init__()\n",
    "        self.tok_emb = TokenEmbedding(vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEncoding(d_model, max_len)\n",
    "        self.drop_out = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        tok_emb = self.tok_emb(x)\n",
    "        pos_emb = self.pos_emb(x)       \n",
    "        return self.drop_out(tok_emb + pos_emb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2e05f345-5473-4d66-b8c8-f9ed5bf08185",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "res: tensor([[[ 1.6016e+00,  0.0000e+00,  7.3131e-01,  ...,  3.4665e-01,\n",
      "           1.2576e-01,  1.7100e+00],\n",
      "         [ 5.4879e-01,  7.7141e-01,  1.6291e+00,  ...,  0.0000e+00,\n",
      "          -0.0000e+00,  2.7418e+00],\n",
      "         [ 0.0000e+00,  1.3316e+00,  0.0000e+00,  ..., -1.0642e+00,\n",
      "          -6.8040e-02, -4.0913e-01]],\n",
      "\n",
      "        [[-1.1485e+00,  3.1954e-01,  0.0000e+00,  ...,  0.0000e+00,\n",
      "           1.0274e+00,  1.3045e+00],\n",
      "         [-5.3425e-02,  2.9367e-01,  5.0048e-01,  ...,  1.2340e+00,\n",
      "           1.0360e+00, -7.7892e-01],\n",
      "         [ 1.0103e+00, -4.6239e-01,  1.0405e+00,  ...,  1.1111e+00,\n",
      "           2.3036e-04,  1.1111e+00]]], grad_fn=<MulBackward0>)\n",
      "res.shape: torch.Size([2, 3, 512])\n"
     ]
    }
   ],
   "source": [
    "te = TransformerEmbedding(1000,512,512,0.1)\n",
    "# x是batch_size为2, seq_len为3\n",
    "x = torch.LongTensor([[6, 5, 4], [3, 2, 1]])\n",
    "res = te.forward(x)\n",
    "print(\"res:\", res)\n",
    "# 返回的形状是[batch_size = 2, seq_len = 3, d_model = 512]\n",
    "print(\"res.shape:\", res.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5ae6f88d-bbd1-4789-9f8f-7b17d9df5d6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScaleDotProductAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    计算单个点积注意力\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super(ScaleDotProductAttention, self).__init__()\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, q, k, v, mask=None, e=1e-12):\n",
    "        # 输入是一个4维的张量\n",
    "        # [batch_size, head, length, d_tensor]\n",
    "        batch_size, head, length, d_tensor = k.size()\n",
    "\n",
    "        # 1.用Key的转置与Query计算点积\n",
    "        k_t = k.transpose(2, 3)  # transpose\n",
    "        score = (q @ k_t) / math.sqrt(d_tensor)  # scaled dot product\n",
    "\n",
    "        # 2.进行掩码，encoder不需要进行掩码，decoder需要进行掩码\n",
    "        if mask is not None:\n",
    "            score = score.masked_fill(mask == 0, -10000)  \n",
    "\n",
    "        # 3.通过softmax使分数范围在[0, 1]之间\n",
    "        score = self.softmax(score)\n",
    "\n",
    "        # 4.再与Value相乘        \n",
    "        v = score @ v\n",
    "\n",
    "        return v, score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c3d594e3-af75-45a0-8558-8174ab9bf1a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, n_head):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.n_head = n_head\n",
    "        self.attention = ScaleDotProductAttention()\n",
    "        self.w_q = nn.Linear(d_model, d_model)\n",
    "        self.w_k = nn.Linear(d_model, d_model)\n",
    "        self.w_v = nn.Linear(d_model, d_model)\n",
    "        self.w_concat = nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        # 1.点积相应的矩阵\n",
    "        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)\n",
    "\n",
    "        # 2.根据头数进行维度拆分\n",
    "        q, k, v = self.split(q), self.split(k), self.split(v)\n",
    "\n",
    "        # 3.进行计算\n",
    "        out, attention = self.attention(q, k, v, mask=mask)\n",
    "\n",
    "        # 4.把拆分的多头再拼起来\n",
    "        out = self.concat(out)\n",
    "        out = self.w_concat(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "    def split(self, tensor):\n",
    "        \"\"\"\n",
    "        根据头数进行维度拆分\n",
    "\n",
    "        :param tensor: [batch_size, length, d_model]\n",
    "        :return: [batch_size, head, length, d_tensor]\n",
    "        \"\"\"\n",
    "        batch_size, length, d_model = tensor.size()\n",
    "\n",
    "        d_tensor = d_model // self.n_head\n",
    "        tensor = tensor.view(batch_size, length, self.n_head, d_tensor).transpose(1, 2)\n",
    "\n",
    "        return tensor\n",
    "\n",
    "    def concat(self, tensor):\n",
    "        \"\"\"\n",
    "        把拆分的多头再拼起来\n",
    "\n",
    "        :param tensor: [batch_size, head, length, d_tensor]\n",
    "        :return: [batch_size, length, d_model]\n",
    "        \"\"\"\n",
    "        batch_size, head, length, d_tensor = tensor.size()\n",
    "        d_model = head * d_tensor\n",
    "\n",
    "        tensor = tensor.transpose(1, 2).contiguous().view(batch_size, length, d_model)\n",
    "        return tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2eeb6cc0-69b8-4ae3-b398-d9fc09bfef4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, hidden, drop_prob=0.1):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.linear1 = nn.Linear(d_model, hidden)\n",
    "        self.linear2 = nn.Linear(hidden, d_model)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.dropout = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.linear2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d91d170e-5bca-4b0b-9a0e-ff289055620c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, d_model, eps=1e-12):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.alpha = nn.Parameter(torch.ones(d_model))\n",
    "        self.bias = nn.Parameter(torch.zeros(d_model))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        var = x.var(-1, unbiased=False, keepdim=True)\n",
    "\n",
    "        out = (x - mean) / torch.sqrt(var + self.eps)\n",
    "        out = self.alpha * out + self.bias\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0f0dc0bf-972d-482e-91ad-538e8d6c4b76",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.attention = MultiHeadAttention(d_model=d_model, n_head=n_head)\n",
    "        self.norm1 = LayerNorm(d_model=d_model)\n",
    "        self.dropout1 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)\n",
    "        self.norm2 = LayerNorm(d_model=d_model)\n",
    "        self.dropout2 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        # 1.计算注意力\n",
    "        _x = x\n",
    "        x = self.attention(q=x, k=x, v=x, mask=src_mask)\n",
    "\n",
    "        # 2.残差连接和层归一化\n",
    "        x = self.dropout1(x)\n",
    "        x = self.norm1(x + _x)\n",
    "\n",
    "        # 3.前馈层\n",
    "        _x = x\n",
    "        x = self.ffn(x)\n",
    "\n",
    "        # 4.最后一次残差连接和层归一化\n",
    "        x = self.dropout2(x)\n",
    "        x = self.norm2(x + _x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e739d3da-ad98-4e77-bcc3-6a1e5e2af421",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "\n",
    "    def __init__(self, enc_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob):\n",
    "        super().__init__()\n",
    "        self.emb = TransformerEmbedding(d_model=d_model,\n",
    "                                        max_len=max_len,\n",
    "                                        vocab_size=enc_voc_size,\n",
    "                                        drop_prob=drop_prob)\n",
    "\n",
    "        self.layers = nn.ModuleList([EncoderLayer(d_model=d_model,\n",
    "                                                  ffn_hidden=ffn_hidden,\n",
    "                                                  n_head=n_head,\n",
    "                                                  drop_prob=drop_prob)\n",
    "                                     for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, x, src_mask):\n",
    "        x = self.emb(x)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, src_mask)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1ae30ffe-e0a6-487e-a035-96bbf2abe09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.self_attention = MultiHeadAttention(d_model=d_model, n_head=n_head)\n",
    "        self.norm1 = LayerNorm(d_model=d_model)\n",
    "        self.dropout1 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.enc_dec_attention = MultiHeadAttention(d_model=d_model, n_head=n_head)\n",
    "        self.norm2 = LayerNorm(d_model=d_model)\n",
    "        self.dropout2 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)\n",
    "        self.norm3 = LayerNorm(d_model=d_model)\n",
    "        self.dropout3 = nn.Dropout(p=drop_prob)\n",
    "\n",
    "    def forward(self, dec, enc, trg_mask, src_mask):    \n",
    "        # 1.对应上面说的第一点\n",
    "        _x = dec\n",
    "        x = self.self_attention(q=dec, k=dec, v=dec, mask=trg_mask)\n",
    "\n",
    "        # 2.残差连接和层归一化\n",
    "        x = self.dropout1(x)\n",
    "        x = self.norm1(x + _x)\n",
    "\n",
    "        if enc is not None:\n",
    "            # 3.对应上面说的第二点\n",
    "            _x = x\n",
    "            x = self.enc_dec_attention(q=x, k=enc, v=enc, mask=src_mask)\n",
    "\n",
    "            # 4.残差连接和层归一化\n",
    "            x = self.dropout2(x)\n",
    "            x = self.norm2(x + _x)\n",
    "\n",
    "        # 5.前馈层\n",
    "        _x = x\n",
    "        x = self.ffn(x)\n",
    "\n",
    "        # 6.残差连接和层归一化\n",
    "        x = self.dropout3(x)\n",
    "        x = self.norm3(x + _x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a462ceab-1c6c-47af-96ee-09c8b7bcd835",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob):\n",
    "        super().__init__()\n",
    "        self.emb = TransformerEmbedding(d_model=d_model,\n",
    "                                        drop_prob=drop_prob,\n",
    "                                        max_len=max_len,\n",
    "                                        vocab_size=dec_voc_size)\n",
    "\n",
    "        self.layers = nn.ModuleList([DecoderLayer(d_model=d_model,\n",
    "                                                  ffn_hidden=ffn_hidden,\n",
    "                                                  n_head=n_head,\n",
    "                                                  drop_prob=drop_prob)\n",
    "                                     for _ in range(n_layers)])\n",
    "\n",
    "        self.linear = nn.Linear(d_model, dec_voc_size)\n",
    "\n",
    "    def forward(self, trg, src, trg_mask, src_mask):\n",
    "        trg = self.emb(trg)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            trg = layer(trg, src, trg_mask, src_mask)\n",
    "\n",
    "        # 最后经过一个全连接层\n",
    "        output = self.linear(trg)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9b6c8832-6e42-457b-a472-913223f9d415",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "\n",
    "    def __init__(self, src_pad_idx, trg_pad_idx, trg_sos_idx, enc_voc_size, dec_voc_size, d_model, n_head, max_len,\n",
    "                 ffn_hidden, n_layers, drop_prob):\n",
    "        super().__init__()\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "        self.trg_sos_idx = trg_sos_idx\n",
    "        self.encoder = Encoder(d_model=d_model,\n",
    "                               n_head=n_head,\n",
    "                               max_len=max_len,\n",
    "                               ffn_hidden=ffn_hidden,\n",
    "                               enc_voc_size=enc_voc_size,\n",
    "                               drop_prob=drop_prob,\n",
    "                               n_layers=n_layers)\n",
    "\n",
    "        self.decoder = Decoder(d_model=d_model,\n",
    "                               n_head=n_head,\n",
    "                               max_len=max_len,\n",
    "                               ffn_hidden=ffn_hidden,\n",
    "                               dec_voc_size=dec_voc_size,\n",
    "                               drop_prob=drop_prob,\n",
    "                               n_layers=n_layers)\n",
    "\n",
    "    def forward(self, src, trg):\n",
    "        src_mask = self.make_src_mask(src)\n",
    "        trg_mask = self.make_trg_mask(trg)\n",
    "        enc_src = self.encoder(src, src_mask)\n",
    "        output = self.decoder(trg, enc_src, trg_mask, src_mask)\n",
    "        return output\n",
    "\n",
    "    def make_src_mask(self, src):\n",
    "        \"\"\"\n",
    "        创建源序列（src）的掩码, 将pad补零的位置设为False\n",
    "        \"\"\"\n",
    "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
    "        return src_mask\n",
    "\n",
    "    def make_trg_mask(self, trg):\n",
    "        \"\"\"\n",
    "        创建目标序列（trg）的掩码, 1.pad补零的位置设为False；\n",
    "        2.创建一个下三角矩阵，这个矩阵的对角线及以下的为为True，其余位置为False\n",
    "        表示在训练时模型只能依赖于当前和过去的信息，不能依赖未来的信息\n",
    "        \"\"\"\n",
    "        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(3)\n",
    "        trg_len = trg.shape[1]\n",
    "        trg_sub_mask = torch.tril(torch.ones(trg_len, trg_len)).type(torch.ByteTensor).to(self.device)\n",
    "        trg_mask = trg_pad_mask & trg_sub_mask\n",
    "        return trg_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "dd10f763-e89a-4668-bf0d-290a9de3a018",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'device' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[19], line 18\u001b[0m\n\u001b[1;32m     15\u001b[0m enc_voc_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m32000\u001b[39m\n\u001b[1;32m     16\u001b[0m dec_voc_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m25000\u001b[39m\n\u001b[0;32m---> 18\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mTransformer\u001b[49m\u001b[43m(\u001b[49m\u001b[43msrc_pad_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msrc_pad_idx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mtrg_pad_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrg_pad_idx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     20\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mtrg_sos_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrg_sos_idx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     21\u001b[0m \u001b[43m                    \u001b[49m\u001b[43md_model\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43md_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     22\u001b[0m \u001b[43m                    \u001b[49m\u001b[43menc_voc_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43menc_voc_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     23\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mdec_voc_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdec_voc_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     24\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mmax_len\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmax_len\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     25\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mffn_hidden\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mffn_hidden\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mn_head\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_heads\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     27\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mn_layers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_layers\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     28\u001b[0m \u001b[43m                    \u001b[49m\u001b[43mdrop_prob\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdrop_prob\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcount_parameters\u001b[39m(model):\n\u001b[1;32m     31\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28msum\u001b[39m(p\u001b[38;5;241m.\u001b[39mnumel() \u001b[38;5;28;01mfor\u001b[39;00m p \u001b[38;5;129;01min\u001b[39;00m model\u001b[38;5;241m.\u001b[39mparameters() \u001b[38;5;28;01mif\u001b[39;00m p\u001b[38;5;241m.\u001b[39mrequires_grad)\n",
      "Cell \u001b[0;32mIn[17], line 9\u001b[0m, in \u001b[0;36mTransformer.__init__\u001b[0;34m(self, src_pad_idx, trg_pad_idx, trg_sos_idx, enc_voc_size, dec_voc_size, d_model, n_head, max_len, ffn_hidden, n_layers, drop_prob)\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrg_pad_idx \u001b[38;5;241m=\u001b[39m trg_pad_idx\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrg_sos_idx \u001b[38;5;241m=\u001b[39m trg_sos_idx\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice \u001b[38;5;241m=\u001b[39m \u001b[43mdevice\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mencoder \u001b[38;5;241m=\u001b[39m Encoder(d_model\u001b[38;5;241m=\u001b[39md_model,\n\u001b[1;32m     11\u001b[0m                        n_head\u001b[38;5;241m=\u001b[39mn_head,\n\u001b[1;32m     12\u001b[0m                        max_len\u001b[38;5;241m=\u001b[39mmax_len,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     15\u001b[0m                        drop_prob\u001b[38;5;241m=\u001b[39mdrop_prob,\n\u001b[1;32m     16\u001b[0m                        n_layers\u001b[38;5;241m=\u001b[39mn_layers)\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdecoder \u001b[38;5;241m=\u001b[39m Decoder(d_model\u001b[38;5;241m=\u001b[39md_model,\n\u001b[1;32m     19\u001b[0m                        n_head\u001b[38;5;241m=\u001b[39mn_head,\n\u001b[1;32m     20\u001b[0m                        max_len\u001b[38;5;241m=\u001b[39mmax_len,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     23\u001b[0m                        drop_prob\u001b[38;5;241m=\u001b[39mdrop_prob,\n\u001b[1;32m     24\u001b[0m                        n_layers\u001b[38;5;241m=\u001b[39mn_layers)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'device' is not defined"
     ]
    }
   ],
   "source": [
    "# 原文base模型参数\n",
    "max_len = 256\n",
    "d_model = 512\n",
    "n_layers = 6\n",
    "n_heads = 8\n",
    "ffn_hidden = 2048\n",
    "drop_prob = 0.1\n",
    "\n",
    "# 分词，词表的一些参数，和通过数据集训练的tokenizer是相关的，这里就简单给一下\n",
    "# 原文使用的数据集是WMT14 EN-DE，enc_voc_size为32000，dec_voc_size为25000，\n",
    "# 这个训练出的词表数值不一样，最终模型的参数也不一样，因为这两个参数会影响Embedding层的参数\n",
    "src_pad_idx = 1\n",
    "trg_pad_idx = 1\n",
    "trg_sos_idx = 2\n",
    "enc_voc_size = 32000\n",
    "dec_voc_size = 25000\n",
    "\n",
    "model = Transformer(src_pad_idx=src_pad_idx,\n",
    "                    trg_pad_idx=trg_pad_idx,\n",
    "                    trg_sos_idx=trg_sos_idx,\n",
    "                    d_model=d_model,\n",
    "                    enc_voc_size=enc_voc_size,\n",
    "                    dec_voc_size=dec_voc_size,\n",
    "                    max_len=max_len,\n",
    "                    ffn_hidden=ffn_hidden,\n",
    "                    n_head=n_heads,\n",
    "                    n_layers=n_layers,\n",
    "                    drop_prob=drop_prob)\n",
    "\n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "# 模型的参数\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')\n",
    "\n",
    "# 模型的结构\n",
    "print(model)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8303a9a-edd4-4683-b31d-58c701937370",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}