examples/gru_lstm.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A rough copy of https://jaketae.github.io/study/pytorch-rnn/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x74043c07bf90>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from torch import nn\n",
    "from data import fetch_names, name_to_array\n",
    "random.seed(1337)\n",
    "torch.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GRUNet(nn.Module):\n",
    "    def __init__(self, input_size: int, num_layers: int, hidden_size: int, output_size: int):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.gru = nn.GRU(input_size, hidden_size, num_layers)\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "    \n",
    "    def __call__(self, x: torch.tensor) -> torch.tensor:\n",
    "        hidden_state = self.init_hidden()\n",
    "        output, _ = self.gru(x, hidden_state)\n",
    "        output = self.fc(output[-1])\n",
    "        return output\n",
    "    \n",
    "    def init_hidden(self) -> torch.tensor:\n",
    "        return torch.zeros(self.num_layers, 1, self.hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set, test_set, char_to_idx, lang_to_label = fetch_names(torch.tensor, torch.tensor)\n",
    "num_letters, num_langs = len(char_to_idx), len(lang_to_label)\n",
    "label_to_lang = {label: lang for lang, label in lang_to_label.items()}\n",
    "\n",
    "model = GRUNet(num_letters, 2, 256, num_langs)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "learning_rate = 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.00: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [01:03<00:00, 285.65it/s]\n",
      "loss 0.93: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [01:07<00:00, 267.62it/s]\n"
     ]
    }
   ],
   "source": [
    "for _ in range(2):\n",
    "    random.shuffle(train_set)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    for i, (name, label) in (t := tqdm(enumerate(train_set), total=len(train_set))):\n",
    "        output = model(name)\n",
    "        loss = criterion(output, label)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        t.set_description(f\"loss {loss.item():.2f}\")\n",
    "    learning_rate /= 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy 0.8226208271051321\n"
     ]
    }
   ],
   "source": [
    "num_correct = 0\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for name, label in test_set:\n",
    "        output = model(name)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "        num_correct += bool(pred == label)\n",
    "print(f\"Accuracy {num_correct / len(test_set)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(name: str) -> str:\n",
    "    model.eval()\n",
    "    name = torch.tensor(name_to_array(name, char_to_idx))\n",
    "    with torch.no_grad():\n",
    "        output = model(name)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "    return label_to_lang[pred.item()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jake: Russian\n",
      "Qin: Chinese\n",
      "Fernando: Italian\n",
      "Demirkan: Russian\n"
     ]
    }
   ],
   "source": [
    "for name in (\"Jake\", \"Qin\", \"Fernando\", \"Demirkan\"):\n",
    "    print(f\"{name}: {predict(name)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTMNet(nn.Module):\n",
    "    def __init__(self, input_size: int, num_layers: int, hidden_size: int, output_size: int):\n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)\n",
    "        self.fc = nn.Linear(hidden_size, output_size)\n",
    "    \n",
    "    def __call__(self, x: torch.tensor) -> torch.tensor:\n",
    "        hidden_state = self.init_hidden()\n",
    "        output, _ = self.lstm(x, hidden_state)\n",
    "        output = self.fc(output[-1])\n",
    "        return output\n",
    "    \n",
    "    def init_hidden(self) -> torch.tensor:\n",
    "        return [torch.zeros(self.num_layers, 1, self.hidden_size) for _ in range(2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.00: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [01:03<00:00, 286.00it/s]\n",
      "loss 0.72: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [00:43<00:00, 412.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy 0.8256103637269556\n",
      "Jake: English\n",
      "Qin: Chinese\n",
      "Fernando: Italian\n",
      "Demirkan: Russian\n"
     ]
    }
   ],
   "source": [
    "model = LSTMNet(num_letters, 2, 256, num_langs)\n",
    "learning_rate = 0.001\n",
    "\n",
    "for _ in range(2):\n",
    "    random.shuffle(train_set)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    for i, (name, label) in (t := tqdm(enumerate(train_set), total=len(train_set))):\n",
    "        output = model(name)\n",
    "        loss = criterion(output, label)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        t.set_description(f\"loss {loss.item():.2f}\")\n",
    "    learning_rate /= 2\n",
    "\n",
    "num_correct = 0\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for name, label in test_set:\n",
    "        output = model(name)\n",
    "        pred = torch.argmax(output, dim=1)\n",
    "        num_correct += bool(pred == label)\n",
    "print(f\"Accuracy {num_correct/len(test_set)}\")\n",
    "\n",
    "for name in (\"Jake\", \"Qin\", \"Fernando\", \"Demirkan\"):\n",
    "    print(f\"{name}: {predict(name)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}