Self-Attention-Machanism.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Embedding an Input Sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Life': 0, 'dessert': 1, 'eat': 2, 'first': 3, 'is': 4, 'short': 5}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentence = 'Life is short, eat dessert first'\n",
    "\n",
    "dc = {w: v for v, w in enumerate(sorted(sentence.replace(',', '').split()))}\n",
    "dc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we use this dictionary to assign an integer index to each word:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 4, 5, 2, 1, 3])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "sentence_int = torch.tensor([dc[w] for w in sentence.replace(',', '').split()])\n",
    "sentence_int"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, using the integer-vector representation of the input sentence, we can use an embedding layer to encode the inputs into a real-vector embedding. Here, we will use a 16-dimensional embedding such that each input word is represented by a 16-dimensional vector. Since the sentence consists of 6 words, this will result in a 6×16-dimensional embedding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 16])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(42)\n",
    "import torch.nn as nn\n",
    "\n",
    "embed = nn.Embedding(10, 16)\n",
    "embedded_sentence = embed(sentence_int).detach()\n",
    "\n",
    "embedded_sentence.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining the Weight Matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "d = embedded_sentence.shape[1]\n",
    "\n",
    "U_query = torch.rand(d, d)\n",
    "U_key = torch.rand(d, d)\n",
    "U_value = torch.rand(d, d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 3, 3])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = torch.rand((2,3,3,4))\n",
    "b = torch.rand((2,3,3,4))\n",
    "c = torch.rand((2,3,3,4))\n",
    "\n",
    "bt = torch.transpose(b, 2, 3)\n",
    "\n",
    "s = a @ bt\n",
    "s.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 3, 3])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "s = F.softmax(s, dim=3)\n",
    "s.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 3, 4])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c = s @ c\n",
    "c.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "\n",
    "        self.n_head = config.n_head\n",
    "        self.dim = config.dim\n",
    "        self.device = device\n",
    "\n",
    "        self.wq = nn.Linear(self.dim, self.dim).to(self.device)\n",
    "        self.wk = nn.Linear(self.dim, self.dim).to(self.device)\n",
    "        self.wv = nn.Linear(self.dim, self.dim).to(self.device)\n",
    "\n",
    "        self.softmax = F.softmax(dim=3)\n",
    "        self.fc = nn.Linear(self.dim, self.dim)\n",
    "        self.norm = nn.LayerNorm(self.dim)\n",
    "        self.dropout = nn.Dropout(config.drop_out)\n",
    "\n",
    "    def split(self, tensor):\n",
    "        a, b, c = tensor.size()\n",
    "        d = c // self.n_head\n",
    "        return tensor.reshape(a, b, self.n_head, d).permute(0, 2, 1, 3) # (a, self.n_head, b, d)\n",
    "\n",
    "    def attention(self, k, q, v):   # dv_model * n_head = d_model\n",
    "        _,_,_,d = k.size()\n",
    "        kt = torch.transpose(k, 2, 3)\n",
    "        # (batch_size, n_head, seq_len, dv_model) @ (batch_size, n_head, dv_model, seq_len) = (batch_size, n_head, seq_len, seq_len)\n",
    "        s = (q @ kt) / math.sqrt(d)   \n",
    "        s = self.somtmax(s)\n",
    "        # (batch_size, n_head, seq_len, seq_len) @ (batch_size, n_head, seq_len, dv_model) = (batch_size, n_head, seq_len, dv_model)\n",
    "        v = s @ v\n",
    "        return v\n",
    "    \n",
    "    def concat(self, tensor):\n",
    "        a, b, c, d = tensor.size()  # (batch_size, n_head, seq_len, dv_model)\n",
    "        tensor = tensor.permute(0, 2, 1, 3) # (batch_size, seq_len, n_head, dv_model)\n",
    "        return tensor.reshape(a, c, b * d)  #  (batch_size, seq_len, n_head * dv_model) => (batch_size, seq_len, d_model) \n",
    "    \n",
    "    def forward(self, x):\n",
    "        # input进去后经过线性变化得到 k,q,v矩阵(batch_size, seq_len, d_model)\n",
    "        k, q, v = self.wk(x), self.wq(x), self.wv(x) \n",
    "\n",
    "        # 将k,q,v分头，分头后的维度为(a, self.n_head, b, d) (batch_size, n_head, seq_len, dv_model)\n",
    "        k, q, v = self.split(k), self.split(q), self.split(v) \n",
    "        \n",
    "        # 经过attention,输出维度为(batch_size, n_head, seq_len, dv_model)\n",
    "        v = self.attention(k, q, v)\n",
    "        \n",
    "        # 将结果concat起来,输出维度为(batch_size, seq_len, d_model)\n",
    "        v = self.concat(v)\n",
    "        \n",
    "        # 备份作残差连接\n",
    "        vb = v \n",
    "\n",
    "        v = self.norm(v + x)\n",
    "\n",
    "        v = self.f(v)\n",
    "\n",
    "        v = self.norm(v + vb)\n",
    "\n",
    "        return self.dropout(v)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 2, 12])\n"
     ]
    }
   ],
   "source": [
    "a = torch.rand((2,3,2,4)) # a2 b3 c2 d4\n",
    "b = torch.rand((2,3,3,4))\n",
    "c = torch.rand((2,3,3,4))\n",
    "s = torch.rand((2,3,3,3))\n",
    "v = torch.rand((2,3,1,3))\n",
    "\n",
    "# a = a.permute(0, 2, 1, 3)\n",
    "print(a.reshape((2,2,3*4)).size())\n",
    "# a.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 ('deeplearn')",
   "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.8.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e03da6500aee53851458a6f61643a708392a6b4f9f0eac272f0a88d0b3352bd9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}