train_with_rnns.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "orig_nbformat": 2,
    "colab": {
      "name": "train_with_rnns.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "GPU",
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ei0eCXMQwnHz"
      },
      "source": [
        "# import libraries\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "from tqdm import tqdm\n",
        "from collections import OrderedDict\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from utils import device, get_num_correct, RunBuilder\n",
        "from rnns import RNN, GRU, LSTM, BLSTM"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7T_jEYxMwnIB"
      },
      "source": [
        "# declare hyperparameters\n",
        "lr = 0.001\n",
        "batch_size = 64\n",
        "input_size = 28\n",
        "hidden_size = 256\n",
        "num_layers = 2\n",
        "num_epochs = 8"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cmIHKE21wnID"
      },
      "source": [
        "# extract and transform the data\n",
        "train_set = torchvision.datasets.MNIST(\n",
        "    root='./data/',\n",
        "    train=True,\n",
        "    download=True,\n",
        "    transform=transforms.ToTensor()\n",
        ")\n",
        "test_set = torchvision.datasets.MNIST(\n",
        "    root='./data/',\n",
        "    train=False,\n",
        "    download=True,\n",
        "    transform=transforms.ToTensor()\n",
        ")\n",
        "# prepare the data loaders\n",
        "train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=1)\n",
        "test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, num_workers=1)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8rldTQXpwnIE"
      },
      "source": [
        "# make an OrderedDict of RNNs\n",
        "runs = OrderedDict(\n",
        "    models = [\n",
        "        RNN(input_size, hidden_size, num_layers),\n",
        "        GRU(input_size, hidden_size, num_layers),\n",
        "        LSTM(input_size, hidden_size, num_layers),\n",
        "        BLSTM(input_size, hidden_size, num_layers)\n",
        "    ]\n",
        ")"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "knSKZ6DkwnIF",
        "outputId": "abbd1669-ca07-43f8-c937-0834df43b4e0"
      },
      "source": [
        "# loss function (categorical cross-entropy)\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "# iterate models in runs and train\n",
        "for run in RunBuilder.get_runs(runs):\n",
        "    model = run.models.to(device)\n",
        "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
        "\n",
        "    net = type(model).__name__\n",
        "    comment = f'-{net}'\n",
        "    tb = SummaryWriter(comment=comment)\n",
        "\n",
        "    for epoch in range(num_epochs):\n",
        "        train_loss, train_correct = 0, 0\n",
        "\n",
        "        model.train()\n",
        "        train_loop = tqdm(train_loader)\n",
        "        for batch in train_loop:\n",
        "            images, labels = batch[0].squeeze(1).to(device), batch[1].to(device)\n",
        "            preds = model(images)\n",
        "            loss = criterion(preds, labels)\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            train_loss += loss.item() * labels.size(0)\n",
        "            train_correct += get_num_correct(preds, labels)\n",
        "\n",
        "            train_loop.set_description(f'{net:6s}Epoch [{epoch+1:2d}/{num_epochs}]')\n",
        "            train_loop.set_postfix(loss=train_loss, acc=train_correct/len(train_set))\n",
        "\n",
        "\n",
        "        tb.add_scalar('Train Loss', train_loss, epoch)\n",
        "        tb.add_scalar('Train Accuracy', train_correct/len(train_set), epoch)\n",
        "        \n",
        "\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            test_loss, test_correct = 0, 0\n",
        "\n",
        "            for batch in test_loader:\n",
        "                images, labels = batch[0].squeeze(1).to(device), batch[1].to(device)\n",
        "                preds = model(images)\n",
        "                loss = criterion(preds, labels)\n",
        "\n",
        "                test_loss += loss.item() * labels.size(0)\n",
        "                test_correct += get_num_correct(preds, labels)\n",
        "\n",
        "            tb.add_scalar('Test Loss', test_loss, epoch)\n",
        "            tb.add_scalar('Test Accuracy', test_correct / len(test_set), epoch)\n",
        "\n",
        "\n",
        "        for name, weight in model.named_parameters():\n",
        "            tb.add_histogram(name, weight, epoch)\n",
        "            tb.add_histogram(f'{name}.grad', weight.grad, epoch)\n",
        "\n",
        "    torch.save(model.state_dict(), f'./models/with_rnns/model{comment}.ckpt')"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "RNN   Epoch [ 1/8]: 100%|██████████| 938/938 [00:13<00:00, 67.41it/s, acc=0.821, loss=3.32e+4]\n",
            "RNN   Epoch [ 2/8]: 100%|██████████| 938/938 [00:13<00:00, 71.10it/s, acc=0.926, loss=1.55e+4]\n",
            "RNN   Epoch [ 3/8]: 100%|██████████| 938/938 [00:13<00:00, 69.91it/s, acc=0.938, loss=1.32e+4]\n",
            "RNN   Epoch [ 4/8]: 100%|██████████| 938/938 [00:13<00:00, 68.04it/s, acc=0.947, loss=1.12e+4]\n",
            "RNN   Epoch [ 5/8]: 100%|██████████| 938/938 [00:13<00:00, 67.02it/s, acc=0.949, loss=1.07e+4]\n",
            "RNN   Epoch [ 6/8]: 100%|██████████| 938/938 [00:13<00:00, 67.76it/s, acc=0.96, loss=8.58e+3]\n",
            "RNN   Epoch [ 7/8]: 100%|██████████| 938/938 [00:13<00:00, 68.67it/s, acc=0.96, loss=8.48e+3]\n",
            "RNN   Epoch [ 8/8]: 100%|██████████| 938/938 [00:13<00:00, 70.39it/s, acc=0.952, loss=1e+4]\n",
            "GRU   Epoch [ 1/8]: 100%|██████████| 938/938 [00:14<00:00, 66.80it/s, acc=0.895, loss=1.89e+4]\n",
            "GRU   Epoch [ 2/8]: 100%|██████████| 938/938 [00:14<00:00, 66.33it/s, acc=0.977, loss=4.38e+3]\n",
            "GRU   Epoch [ 3/8]: 100%|██████████| 938/938 [00:14<00:00, 65.22it/s, acc=0.985, loss=2.95e+3]\n",
            "GRU   Epoch [ 4/8]: 100%|██████████| 938/938 [00:14<00:00, 66.04it/s, acc=0.987, loss=2.32e+3]\n",
            "GRU   Epoch [ 5/8]: 100%|██████████| 938/938 [00:14<00:00, 66.63it/s, acc=0.99, loss=1.84e+3]\n",
            "GRU   Epoch [ 6/8]: 100%|██████████| 938/938 [00:13<00:00, 67.68it/s, acc=0.992, loss=1.58e+3]\n",
            "GRU   Epoch [ 7/8]: 100%|██████████| 938/938 [00:14<00:00, 64.11it/s, acc=0.992, loss=1.46e+3]\n",
            "GRU   Epoch [ 8/8]: 100%|██████████| 938/938 [00:14<00:00, 63.83it/s, acc=0.994, loss=1.12e+3]\n",
            "LSTM  Epoch [ 1/8]: 100%|██████████| 938/938 [00:14<00:00, 65.78it/s, acc=0.868, loss=2.37e+4]\n",
            "LSTM  Epoch [ 2/8]: 100%|██████████| 938/938 [00:14<00:00, 65.53it/s, acc=0.971, loss=5.63e+3]\n",
            "LSTM  Epoch [ 3/8]: 100%|██████████| 938/938 [00:13<00:00, 67.45it/s, acc=0.98, loss=3.74e+3]\n",
            "LSTM  Epoch [ 4/8]: 100%|██████████| 938/938 [00:14<00:00, 64.49it/s, acc=0.986, loss=2.84e+3]\n",
            "LSTM  Epoch [ 5/8]: 100%|██████████| 938/938 [00:14<00:00, 64.09it/s, acc=0.988, loss=2.45e+3]\n",
            "LSTM  Epoch [ 6/8]: 100%|██████████| 938/938 [00:14<00:00, 65.81it/s, acc=0.989, loss=2.06e+3]\n",
            "LSTM  Epoch [ 7/8]: 100%|██████████| 938/938 [00:14<00:00, 66.65it/s, acc=0.991, loss=1.7e+3]\n",
            "LSTM  Epoch [ 8/8]: 100%|██████████| 938/938 [00:14<00:00, 66.29it/s, acc=0.992, loss=1.49e+3]\n",
            "BLSTM Epoch [ 1/8]: 100%|██████████| 938/938 [00:18<00:00, 50.14it/s, acc=0.855, loss=2.6e+4]\n",
            "BLSTM Epoch [ 2/8]: 100%|██████████| 938/938 [00:18<00:00, 50.29it/s, acc=0.97, loss=6.19e+3]\n",
            "BLSTM Epoch [ 3/8]: 100%|██████████| 938/938 [00:17<00:00, 52.22it/s, acc=0.98, loss=3.93e+3]\n",
            "BLSTM Epoch [ 4/8]: 100%|██████████| 938/938 [00:18<00:00, 51.48it/s, acc=0.985, loss=3.04e+3]\n",
            "BLSTM Epoch [ 5/8]: 100%|██████████| 938/938 [00:18<00:00, 50.79it/s, acc=0.987, loss=2.53e+3]\n",
            "BLSTM Epoch [ 6/8]: 100%|██████████| 938/938 [00:18<00:00, 52.07it/s, acc=0.989, loss=2.09e+3]\n",
            "BLSTM Epoch [ 7/8]: 100%|██████████| 938/938 [00:18<00:00, 51.81it/s, acc=0.99, loss=1.93e+3]\n",
            "BLSTM Epoch [ 8/8]: 100%|██████████| 938/938 [00:18<00:00, 51.68it/s, acc=0.992, loss=1.67e+3]\n"
          ],
          "name": "stderr"
        }
      ]
    }
  ]
}