pos_enc_and_multiple_layers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "sqH2SEO-X9T6"
      },
      "cell_type": "markdown",
      "source": [
        "# Neural Networks on MNIST\n",
        "\n",
        "This Jupyter notebook explains various approaches for implementing neural networks that recognize digits on [MNIST](http://yann.lecun.com/exdb/mnist/) dataset."
      ]
    },
    {
      "metadata": {
        "id": "OaLeoEcDYVO_"
      },
      "cell_type": "markdown",
      "source": [
        "## Preparing the MNIST dataset\n",
        "\n",
        "Most deep learning frameworks provide APIs for loading famous datasets like MNIST (e.g., `torchvision.datasets.MNIST` in pytorch). The APIs are handy, but hide the important step for preparing a training data for a deep learning framework; when graduating from an example dataset to the real data, we must convert a training data of our interest into the data structure that is acceptable by a deep learning framework.\n",
        "\n",
        "The cell below downloads the original distribution of the MNIST dataset on the Web, converts the dataset into `numpy` arrays, and saves the arrays as the file `mnist.npz` with keyword names."
      ]
    },
    {
      "metadata": {
        "id": "OMaOcZMBPuQY"
      },
      "cell_type": "code",
      "source": [
        "import gzip\n",
        "import os\n",
        "import sys\n",
        "import struct\n",
        "import numpy as np\n",
        "\n",
        "def read_image(fi):\n",
        "    magic, n, rows, columns = struct.unpack(\">IIII\", fi.read(16))\n",
        "    assert magic == 0x00000803\n",
        "    assert rows == 28\n",
        "    assert columns == 28\n",
        "    rawbuffer = fi.read()\n",
        "    assert len(rawbuffer) == n * rows * columns\n",
        "    rawdata = np.frombuffer(rawbuffer, dtype='>u1', count=n*rows*columns)\n",
        "    return rawdata.reshape(n, rows, columns).astype(np.float32) / 255.0\n",
        "\n",
        "def read_label(fi):\n",
        "    magic, n = struct.unpack(\">II\", fi.read(8))\n",
        "    assert magic == 0x00000801\n",
        "    rawbuffer = fi.read()\n",
        "    assert len(rawbuffer) == n\n",
        "    return np.frombuffer(rawbuffer, dtype='>u1', count=n)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    os.system('wget -N http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz')\n",
        "    os.system('wget -N http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz')\n",
        "    os.system('wget -N http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz')\n",
        "    os.system('wget -N http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz')\n",
        "\n",
        "    np.savez_compressed(\n",
        "        'mnist',\n",
        "        train_x=read_image(gzip.open('train-images-idx3-ubyte.gz', 'rb')),\n",
        "        train_y=read_label(gzip.open('train-labels-idx1-ubyte.gz', 'rb')),\n",
        "        test_x=read_image(gzip.open('t10k-images-idx3-ubyte.gz', 'rb')),\n",
        "        test_y=read_label(gzip.open('t10k-labels-idx1-ubyte.gz', 'rb'))\n",
        "    )"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "metadata": {
        "id": "TStlGwaUaZKC"
      },
      "cell_type": "markdown",
      "source": [
        "The file contains four numpy arrays (one tensor and array for each split of training and test sets) with the keywords:\n",
        "\n",
        "+ `train_x`: $60000 \\text{ (images)} \\times 28 \\text{ (y)} \\times 28 \\text{ (x)}$\n",
        "+ `train_y`: $60000 \\text{ (labels)}$\n",
        "+ `test_x`: $10000 \\text{ (images)} \\times 28 \\text{ (y)} \\times 28 \\text{ (x)}$\n",
        "+ `test_y`: $10000 \\text{ (labels)}$\n"
      ]
    },
    {
      "metadata": {
        "id": "QwF495MIe0-e"
      },
      "cell_type": "markdown",
      "source": [
        "### Install pytorch"
      ]
    },
    {
      "metadata": {
        "id": "h8U839iGXCH3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "600dff09-c0ce-4a34-88b4-27b05fee529b"
      },
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (2.0.1+cu118)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (0.15.2+cu118)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch) (3.12.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch) (3.1.2)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch) (2.0.0)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch) (3.27.4.1)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch) (16.0.6)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torchvision) (1.23.5)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from torchvision) (2.31.0)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision) (9.4.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch) (2.1.3)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (3.2.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (3.4)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (2.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->torchvision) (2023.7.22)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch) (1.3.0)\n"
          ]
        }
      ]
    },
    {
      "metadata": {
        "id": "kjBd7K4lI3W5"
      },
      "cell_type": "markdown",
      "source": [
        "### Transformers"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "data = \"Transformers is a media franchise produced by American toy company Hasbro and Japanese toy company Takara Tomy   It primarily follows the heroic Autobots and the villainous Decepticons   two alien robot factions at war that can transform into other forms   such as vehicles and animals   The franchise encompasses toys   animation   comic books   video games and films   As of            it generated more than      trillion          billion   in revenue         making it one of the highest  grossing media franchises of all time  The franchise began in          with the Transformers toy line   comprising transforming mecha toys from Takara  s Diaclone and Micro Change toylines rebranded for Western markets         The term   Generation      covers both the animated television series The Transformers and the comic book series of the same name   which are further divided into Japanese   British and Canadian spin  offs   respectively   Sequels followed   such as the Generation    comic book and Beast Wars TV series   which became its own mini  universe   Generation    characters underwent two reboots with Dreamwave Productions in          and IDW Publishing in            with a third starting in            There have been other incarnations of the story based on different toy lines during and after the     th century   The first was the Robots in Disguise series   followed by three shows   Armada   Energon   and Cybertron   that constitute a single universe called the   Unicron Trilogy    A live  action film series started in            again distinct from previous incarnations   while the Transformers   Animated series merged concepts from the G   continuity   the          live  action film and the   Unicron Trilogy     For most of the         s   in an attempt to mitigate the wave of reboots   the   Aligned Continuity   was established   In            Transformers   Cyberverse debuted   once again   distinct from the previous incarnations  Although initially a separate and competing franchise started in            Tonka  s GoBots became the intellectual property of Hasbro after their buyout of Tonka in            Subsequently   the universe depicted in the animated series Challenge of the GoBots and follow  up film GoBots   Battle of the Rock Lords was retroactively established as an alternate universe within the Transformers multiverse        FictionTransformers   Generation                          Main articles   Transformers   Generation    and The Transformers   TV series  Classic Transformers franchise logo used until         Spider  Man battles Megatron on the cover of The Transformers       Generation One is a retroactive term for the Transformers characters that appeared between          and            The Transformers began with the         s Japanese toy lines Micro Change and Diaclone   They presented robots able to transform into everyday vehicles   electronic items or weapons   Hasbro bought the Micro Change and Diaclone toys   and partnered with Takara         Marvel Comics was hired by Hasbro to create the backstory   editor  in  chief Jim Shooter wrote an overall story   and gave the task of creating the characters to writer Dennis O  Neil         Unhappy with O  Neil  s work   although O  Neil created the name   Optimus Prime       Shooter chose Bob Budiansky to create the characters        The Transformers mecha were largely designed by Shōji Kawamori   the creator of the Japanese mecha anime franchise Macross   which was adapted into the Robotech franchise in North America           Kawamori came up with the idea of transforming mechs while working on the Diaclone and Macross franchises in the early         s   such as the VF     Valkyrie in Macross and Robotech     with his Diaclone mechs later providing the basis for Transformers        The primary concept of Generation One is that the heroic Optimus Prime   the villainous Megatron   and their finest soldiers crash land on pre  historic Earth in the Ark and the Nemesis before awakening in            Cybertron hurtling through the Neutral zone as an effect of the war   The Marvel comic was originally part of the main Marvel Universe   with appearances from Spider  Man and Nick Fury   plus some cameos         as well as a visit to the Savage Land          The Transformers TV series began around the same time   Produced by Sunbow Productions and Marvel Productions   later Hasbro Productions   from the start it contradicted Budiansky  s backstories   The TV series shows the Autobots looking for new energy sources   and crash landing as the Decepticons attack           Marvel interpreted the Autobots as destroying a rogue asteroid approaching Cybertron           Shockwave is loyal to Megatron in the TV series   keeping Cybertron in a stalemate during his absence           but in the comic book he attempts to take command of the Decepticons           The TV series would also differ wildly from the origins Budiansky had created for the Dinobots                   the Decepticon turned Autobot Jetfire           known as Skyfire on TV             the Constructicons   who combine to form Devastator                     and Omega Supreme                   The Marvel comic establishes early on that Prime wields the Creation Matrix   which gives life to machines   In the second season   the two  part episode The Key to Vector Sigma introduced the ancient Vector Sigma computer   which served the same original purpose as the Creation Matrix   giving life to Transformers     and its guardian Alpha Trion  In            the cartoon became the film The Transformers   The Movie   which is set in the year            It introduced the Matrix as the   Autobot Matrix of Leadership     as a fatally wounded Prime gives it to Ultra Magnus   however   as Prime dies he drops the matrix   which is then caught by Hot Rod who subsequently becomes Rodimus Prime later on in the film   Unicron   a transformer who devours planets   fears its power and recreates a heavily damaged Megatron as Galvatron   as well as Bombshell or Skywarp becoming Cyclonus   Thundercracker becoming Scourge and two other Insecticons becoming Scourge  s huntsmen   the Sweeps   Eventually   Rodimus Prime takes out the Matrix and destroys Unicron           In the United Kingdom   the weekly comic book interspliced original material to keep up with U  S   reprints           and The Movie provided much new material   Writer Simon Furman proceeded to expand the continuity with movie spin  offs involving the time travelling Galvatron                   The Movie also featured guest voices from Leonard Nimoy as Galvatron   Scatman Crothers as Jazz   Casey Kasem as Cliffjumper   Orson Welles as Unicron and Eric Idle as the leader of the Junkions   Wreck  Gar   though unnamed in the movie     The Transformers theme tune for the film was performed by Lion with   Weird Al   Yankovic adding a song to the soundtrack  The third season followed up The Movie   with the revelation of the Quintessons having used Cybertron as a factory   Their robots rebel   and in time the workers become the Autobots and the soldiers become the Decepticons     Note   This appears to contradict background presented in the first two seasons of the series     It is the Autobots who develop transformation           Due to popular demand           Optimus Prime is resurrected at the conclusion of the third season           and the series ended with a three  episode story arc   However   the Japanese broadcast of the series was supplemented with a newly produced OVA   Scramble City   before creating entirely new series to continue the storyline   ignoring the          end of the American series   The extended Japanese run consisted of The Headmasters   Super  God Masterforce   Victory and Zone   then in illustrated magazine form as Battlestars   Return of Convoy and Operation   Combination   Just as the TV series was wrapping up   Marvel continued to expand its continuity   It followed The Movie  s example by killing Prime         and Megatron           albeit in the present day   Dinobot leader Grimlock takes over as Autobot leader           There was a G  I   Joe crossover         and the limited series The Transformers   Headmasters   which further expanded the scope to the planet Nebulon           It led on to the main title resurrecting Prime as a Powermaster          \"\n",
        "\n",
        "\n",
        "#\"The franchise began in 1984 with the Transformers toy line, comprising transforming mecha toys from Takara's Diaclone and Micro Change toylines rebranded for Western markets.[2] The term Generation 1 covers both the animated television series The Transformers and the comic book series of the same name, which are further divided into Japanese, British and Canadian spin-offs, respectively. Sequels followed, such as the Generation 2 comic book and Beast Wars TV series, which became its own mini-universe. Generation 1 characters underwent two reboots with Dreamwave Productions in 2001 and IDW Publishing in 2005, with a third starting in 2019. There have been other incarnations of the story based on different toy lines during and after the 20th century. The first was the Robots in Disguise series, followed by three shows (Armada, Energon, and Cybertron) that constitute a single universe called the Unicron Trilogy\"\n",
        "\n",
        "data = data.replace(\".\", \"\")\n",
        "data.replace(\",\",\"\")\n",
        "\n",
        "data = data.lower()\n",
        "data = data.split(\" \")\n",
        "\n",
        "\n",
        "uniques = len(set(data))\n",
        "print(uniques)\n",
        "words = list(set(data))\n",
        "\n",
        "numdata = []\n",
        "for word in data:\n",
        "  numdata.append(words.index(word))\n",
        "\n",
        "numdata = torch.IntTensor(numdata)\n",
        "print(numdata)"
      ],
      "metadata": {
        "id": "dSW1Tas9HKyi",
        "outputId": "379abd4a-8357-43ec-f50c-dc7ae6e43bbd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "519\n",
            "tensor([ 10, 230, 269,  ...,   0,   0,   0], dtype=torch.int32)\n"
          ]
        }
      ]
    },
    {
      "metadata": {
        "id": "DdSGWuLffzEg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "72796e67-484a-453e-ad88-2d690429773f"
      },
      "cell_type": "code",
      "source": [
        "\n",
        "import math\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "class self_attention(nn.Module):\n",
        "  def __init__(self, device=\"cuda:0\"):\n",
        "    super().__init__()\n",
        "    self.num_words = 12\n",
        "    self.num_dims = 12\n",
        "\n",
        "    self.pos_encoding = torch.zeros((self.num_words,self.num_dims))\n",
        "\n",
        "    for pos in range(self.num_words):\n",
        "      for i in range(self.num_dims):\n",
        "        if pos % 2 == 0:\n",
        "          self.pos_encoding[pos,i] = math.sin(pos/math.pow(10000, i/self.num_dims))\n",
        "        else:\n",
        "          self.pos_encoding[pos,i] = math.cos(pos/math.pow(10000, i-1/self.num_dims))\n",
        "    self.pos_encoding = self.pos_encoding.to(device)\n",
        "\n",
        "    self.Q = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.K = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.V = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.softmax = nn.Softmax(dim=0)\n",
        "    self.embedding = nn.Embedding(519,self.num_dims).to(device)\n",
        "    self.flatten = nn.Flatten(start_dim=0)\n",
        "    self.preoutput = nn.Linear(self.num_words*self.num_dims,512).to(device)\n",
        "    self.gelu = nn.GELU()\n",
        "    self.output = nn.Linear(512, 519).to(device)\n",
        "\n",
        "  def forward(self, x, device=\"cuda:0\"):\n",
        "    x = x.to(device)\n",
        "    x = self.embedding(x)\n",
        "    x += self.pos_encoding\n",
        "    x_q = self.Q(x)\n",
        "    x_k = self.K(x)\n",
        "    x_v = self.V(x)\n",
        "    x_qk =  x_q.matmul(x_k.transpose(0,1))/math.sqrt(self.num_dims)\n",
        "    x = self.softmax(x_qk)*x_v\n",
        "    x = self.flatten(x)\n",
        "    x=self.preoutput(x)\n",
        "    x = self.gelu(x)\n",
        "    x = self.output(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "torch.manual_seed(0)\n",
        "self_atten = self_attention()\n",
        "\n",
        "#p = self_atten(numdata[:12])\n",
        "\n",
        "optimizer = torch.optim.Adam(self_atten.parameters(), lr = 0.0001)\n",
        "loss = nn.CrossEntropyLoss()\n",
        "\n",
        "for k in range(50):\n",
        "  loss_n = 0\n",
        "  for i in range(len(numdata)-13):\n",
        "    optimizer.zero_grad()\n",
        "    x = numdata[i:i+12]\n",
        "    y = torch.zeros(519)\n",
        "    y = y.to(\"cuda:0\")\n",
        "    y[numdata[i+12]] = 1\n",
        "    y_p = self_atten(x)\n",
        "    l = loss(y_p, y)\n",
        "    l.backward()\n",
        "    optimizer.step()\n",
        "    loss_n += l\n",
        "  loss_n /= len(numdata)-13\n",
        "  print(\"epoch:\", k ,\"loss:\", loss_n)\n",
        "\n",
        "\n",
        "#p = self_atten(numdata[:12])\n",
        "\n",
        "#print(p, words[p.argmax()], data[13])\n",
        "\n",
        "accuracy = 0\n",
        "for i in range(len(numdata)-13):\n",
        "  x = numdata[i:i+12]\n",
        "  y = data[i+12]\n",
        "  y_p = self_atten(x)\n",
        "  y_p = y_p.to(\"cpu\")\n",
        "  #print(y, words[y_p.argmax()])\n",
        "  if y != words[y_p.argmax()]:\n",
        "    print(\"Wrong!\")\n",
        "  else:\n",
        "    accuracy += 1\n",
        "accuracy /= len(numdata)-13\n",
        "print(\"accuracy:\", accuracy)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch: 0 loss: tensor(4.4186, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 1 loss: tensor(3.6963, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 2 loss: tensor(3.6020, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 3 loss: tensor(3.5178, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 4 loss: tensor(3.4197, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 5 loss: tensor(3.2865, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 6 loss: tensor(3.1224, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 7 loss: tensor(2.9508, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 8 loss: tensor(2.7855, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 9 loss: tensor(2.6288, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 10 loss: tensor(2.4811, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 11 loss: tensor(2.3421, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 12 loss: tensor(2.2102, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 13 loss: tensor(2.0828, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 14 loss: tensor(1.9582, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 15 loss: tensor(1.8359, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 16 loss: tensor(1.7163, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 17 loss: tensor(1.6000, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 18 loss: tensor(1.4883, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 19 loss: tensor(1.3821, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 20 loss: tensor(1.2822, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 21 loss: tensor(1.1888, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 22 loss: tensor(1.1021, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 23 loss: tensor(1.0220, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 24 loss: tensor(0.9480, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 25 loss: tensor(0.8797, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 26 loss: tensor(0.8168, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 27 loss: tensor(0.7588, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 28 loss: tensor(0.7053, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 29 loss: tensor(0.6560, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 30 loss: tensor(0.6105, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 31 loss: tensor(0.5685, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 32 loss: tensor(0.5298, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 33 loss: tensor(0.4942, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 34 loss: tensor(0.4611, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 35 loss: tensor(0.4304, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 36 loss: tensor(0.4018, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 37 loss: tensor(0.3750, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 38 loss: tensor(0.3499, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 39 loss: tensor(0.3267, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 40 loss: tensor(0.3051, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 41 loss: tensor(0.2852, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 42 loss: tensor(0.2669, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 43 loss: tensor(0.2500, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 44 loss: tensor(0.2347, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 45 loss: tensor(0.2210, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 46 loss: tensor(0.2086, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 47 loss: tensor(0.1975, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 48 loss: tensor(0.1874, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 49 loss: tensor(0.1783, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "accuracy: 0.9440993788819876\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "class self_attention_pure(nn.Module):\n",
        "  def __init__(self, device=\"cuda:0\", p_in=12, mid = 512, dims=0):\n",
        "    super().__init__()\n",
        "    print(p_in)\n",
        "    self.num_words = p_in\n",
        "    self.num_dims = p_in\n",
        "    if dims != 0:\n",
        "      self.num_dims = dims\n",
        "\n",
        "    self.Q = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.K = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.V = nn.Linear(self.num_dims, self.num_dims).to(device)\n",
        "    self.softmax = nn.Softmax(dim=0)\n",
        "    self.preoutput_0 = nn.Linear(self.num_dims,mid).to(device)\n",
        "    self.preoutput_1 = nn.Linear(mid, self.num_dims).to(device)\n",
        "    self.gelu = nn.GELU()\n",
        "\n",
        "  def forward(self, x, device=\"cuda:0\"):\n",
        "    x = x.to(device)\n",
        "    id = x\n",
        "    x_q = self.Q(x)\n",
        "    x_k = self.K(x)\n",
        "    x_v = self.V(x)\n",
        "    x_qk = x_q*x_k/math.sqrt(self.num_dims)\n",
        "    x = self.softmax(x_qk)*(x_v)\n",
        "    x = self.preoutput_0(x)\n",
        "    x = self.preoutput_1(x)\n",
        "    x += id\n",
        "    x = self.gelu(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "\n",
        "class self_attention_w_pos_enc(nn.Module):\n",
        "  def __init__(self, device=\"cuda:0\", words = 12, dims = 12):\n",
        "    super().__init__()\n",
        "    self.num_words = words\n",
        "    self.num_dims = dims\n",
        "    self.embedding = nn.Embedding(519,self.num_dims).to(device)\n",
        "    self.pos_encoding = torch.zeros((self.num_words,self.num_dims))\n",
        "\n",
        "    for i in range(self.num_words):\n",
        "      for j in range(self.num_dims):\n",
        "        if j % 2 == 0:\n",
        "          self.pos_encoding[i,j] = math.sin(i/math.pow(10000, j/self.num_dims))\n",
        "        else:\n",
        "          self.pos_encoding[i,j] = math.cos(i/math.pow(10000, j-1/self.num_dims))\n",
        "    self.pos_encoding = self.pos_encoding.to(device)\n",
        "\n",
        "    self.attention = self_attention_pure(device=device, p_in = self.num_dims, dims = self.num_dims, mid=512)\n",
        "    self.output = nn.Linear(512, 519).to(device)\n",
        "\n",
        "  def forward(self, x, device=\"cuda:0\"):\n",
        "    x = x.to(device)\n",
        "    x = self.embedding(x)\n",
        "    x += self.pos_encoding\n",
        "    x = self.attention(x)\n",
        "    x = self.output(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "class self_attention_w_pos_enc_pure(nn.Module):\n",
        "  def __init__(self, device=\"cuda:0\", num_words = 12, num_dims=12, out=512):\n",
        "    super().__init__()\n",
        "    self.num_words = num_words\n",
        "    self.num_dims = num_dims\n",
        "    self.embedding = nn.Embedding(519,self.num_dims).to(device)\n",
        "    self.pos_encoding = torch.zeros((self.num_words,self.num_dims))\n",
        "\n",
        "    for i in range(self.num_words):\n",
        "      for j in range(self.num_dims):\n",
        "        if j % 2 == 0:\n",
        "          self.pos_encoding[i,j] = math.sin(i/math.pow(10000, j/self.num_dims))\n",
        "        else:\n",
        "          self.pos_encoding[i,j] = math.cos(i/math.pow(10000, j-1/self.num_dims))\n",
        "    self.pos_encoding = self.pos_encoding.to(device)\n",
        "\n",
        "    self.attention = self_attention_pure(device=device, p_in = self.num_dims, dims = self.num_dims, mid=out)\n",
        "\n",
        "  def forward(self, x, device=\"cuda:0\"):\n",
        "    x = x.to(device)\n",
        "    x = self.embedding(x)\n",
        "    x += self.pos_encoding\n",
        "    x = self.attention(x)\n",
        "    return x\n",
        "\n",
        "\n",
        "class multi_layer_self_attention(nn.Module):\n",
        "  def __init__(self, device=\"cuda:0\", outs=[144,256,512,519], seq_length=12):\n",
        "    super().__init__()\n",
        "    self.attention_w_enc = self_attention_w_pos_enc_pure(device=device, num_words=seq_length, num_dims=12, out=outs[0])\n",
        "    self.attention_1 = self_attention_pure(device=device, p_in=12, mid=outs[1])\n",
        "    self.linear_1 = nn.Linear(seq_length*12, outs[0]).to(device)\n",
        "    self.linear_2 = nn.Linear(outs[0], outs[1]).to(device)\n",
        "    self.flatten = nn.Flatten(start_dim=0)\n",
        "    self.output = nn.Linear(outs[1], outs[3]).to(device)\n",
        "\n",
        "\n",
        "  def forward(self, x, device=\"cuda:0\"):\n",
        "    x = x.to(device)\n",
        "    x = self.attention_w_enc(x)\n",
        "    x = self.attention_1(x)\n",
        "    x = self.flatten(x)\n",
        "    x = self.linear_1(x)\n",
        "    x = self.linear_2(x)\n",
        "    x = self.output(x)\n",
        "    return x\n",
        "\n",
        "torch.manual_seed(0)\n",
        "seq_length = 24\n",
        "self_atten = multi_layer_self_attention(seq_length=seq_length)\n",
        "#p = self_atten(numdata[:12])\n",
        "\n",
        "optimizer = torch.optim.Adam(self_atten.parameters(), lr = 0.0001)\n",
        "loss = nn.CrossEntropyLoss()\n",
        "schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2, cooldown=3, factor=0.333)\n",
        "\n",
        "\n",
        "\n",
        "for k in range(50):\n",
        "  loss_n = 0\n",
        "  for i in range(len(numdata)-(seq_length+1)):\n",
        "    optimizer.zero_grad()\n",
        "    x = numdata[i:i+seq_length]\n",
        "    y = torch.zeros(519)\n",
        "    y = y.to(\"cuda:0\")\n",
        "    y[numdata[i+seq_length]] = 1\n",
        "    y_p = self_atten(x)\n",
        "    l = loss(y_p, y)\n",
        "    l.backward()\n",
        "    optimizer.step()\n",
        "    loss_n += l\n",
        "  loss_n /= len(numdata)-seq_length-1\n",
        "  schedule.step(loss_n)\n",
        "  print(\"epoch:\", k ,\"loss:\", loss_n)\n",
        "\n",
        "\n",
        "#p = self_atten(numdata[:12])\n",
        "\n",
        "#print(p, words[p.argmax()], data[13])\n",
        "\n",
        "accuracy = 0\n",
        "for i in range(len(numdata)-(seq_length+1)):\n",
        "  x = numdata[i:i+seq_length]\n",
        "  y = data[i+seq_length]\n",
        "  y_p = self_atten(x)\n",
        "  y_p = y_p.to(\"cpu\")\n",
        "  #print(y, words[y_p.argmax()])\n",
        "  if y != words[y_p.argmax()]:\n",
        "    print(\"Wrong!\")\n",
        "  else:\n",
        "    accuracy += 1\n",
        "accuracy /= len(numdata)-seq_length-1\n",
        "print(\"accuracy:\", accuracy)"
      ],
      "metadata": {
        "id": "X8H7JxBoLXeh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f0a4d68a-60c6-43cf-fd46-ee2f07e63d06"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "12\n",
            "12\n",
            "epoch: 0 loss: tensor(4.2351, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 1 loss: tensor(3.5792, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 2 loss: tensor(3.3300, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 3 loss: tensor(3.0878, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 4 loss: tensor(2.7788, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 5 loss: tensor(2.3964, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 6 loss: tensor(2.0140, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 7 loss: tensor(1.6950, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 8 loss: tensor(1.4465, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 9 loss: tensor(1.2531, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 10 loss: tensor(1.0971, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 11 loss: tensor(0.9650, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 12 loss: tensor(0.8525, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 13 loss: tensor(0.7556, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 14 loss: tensor(0.6704, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 15 loss: tensor(0.5965, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 16 loss: tensor(0.5318, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 17 loss: tensor(0.4739, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 18 loss: tensor(0.4226, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 19 loss: tensor(0.3767, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 20 loss: tensor(0.3347, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 21 loss: tensor(0.2953, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 22 loss: tensor(0.2583, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 23 loss: tensor(0.2241, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 24 loss: tensor(0.1932, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 25 loss: tensor(0.1653, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 26 loss: tensor(0.1492, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 27 loss: tensor(0.1439, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 28 loss: tensor(0.1284, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 29 loss: tensor(0.1208, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 30 loss: tensor(0.1167, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 31 loss: tensor(0.1021, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 32 loss: tensor(0.0900, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 33 loss: tensor(0.0916, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 34 loss: tensor(0.0886, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 35 loss: tensor(0.0645, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 36 loss: tensor(0.0621, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 37 loss: tensor(0.0773, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 38 loss: tensor(0.0437, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 39 loss: tensor(0.0594, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 40 loss: tensor(0.0480, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 41 loss: tensor(0.0448, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 42 loss: tensor(0.0308, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 43 loss: tensor(0.0149, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 44 loss: tensor(0.0109, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 45 loss: tensor(0.0084, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 46 loss: tensor(0.0069, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 47 loss: tensor(0.0056, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 48 loss: tensor(0.0050, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "epoch: 49 loss: tensor(0.0041, device='cuda:0', grad_fn=<DivBackward0>)\n",
            "Wrong!\n",
            "Wrong!\n",
            "Wrong!\n",
            "accuracy: 0.9985583853916387\n"
          ]
        }
      ]
    }
  ]
}