pokemon_classifier_train.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "pokemon_classifier_train.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kmIOkV65f5S0"
      },
      "source": [
        "### Import Packages\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AU_uqfXvD_ew"
      },
      "source": [
        "# Load the TensorBoard notebook extension\n",
        "%load_ext tensorboard\n",
        "from tensorflow import summary\n",
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RkNT79b8xtbg"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X1f7bWu9lJtX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7d298f28-a4f3-4319-ed5e-21a8ffa5fcc1"
      },
      "source": [
        "import os\n",
        "import csv\n",
        "import numpy as np\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Module for Google Drive\n",
        "from google.colab import drive\n",
        "\n",
        "# Module for Importing Images\n",
        "from PIL import Image \n",
        "\n",
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "print(torch.__version__)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "1.7.0+cu101\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kA3MZRzNHclY"
      },
      "source": [
        "### Import your drive's contents!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JXPZ3dc6W3kp",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 562
        },
        "outputId": "0a5406c3-1e35-4f95-fe95-15c6f3ac5e89"
      },
      "source": [
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36m_input_request\u001b[0;34m(self, prompt, ident, parent, password)\u001b[0m\n\u001b[1;32m    728\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 729\u001b[0;31m                 \u001b[0mident\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreply\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msession\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstdin_socket\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    730\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/jupyter_client/session.py\u001b[0m in \u001b[0;36mrecv\u001b[0;34m(self, socket, mode, content, copy)\u001b[0m\n\u001b[1;32m    802\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 803\u001b[0;31m             \u001b[0mmsg_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msocket\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv_multipart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    804\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mzmq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mZMQError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/zmq/sugar/socket.py\u001b[0m in \u001b[0;36mrecv_multipart\u001b[0;34m(self, flags, copy, track)\u001b[0m\n\u001b[1;32m    565\u001b[0m         \"\"\"\n\u001b[0;32m--> 566\u001b[0;31m         \u001b[0mparts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mflags\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrack\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrack\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    567\u001b[0m         \u001b[0;31m# have first part already, only loop while more to receive\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket.Socket.recv\u001b[0;34m()\u001b[0m\n",
            "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket.Socket.recv\u001b[0;34m()\u001b[0m\n",
            "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket._recv_copy\u001b[0;34m()\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/zmq/backend/cython/checkrc.pxd\u001b[0m in \u001b[0;36mzmq.backend.cython.checkrc._check_rc\u001b[0;34m()\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: ",
            "\nDuring handling of the above exception, another exception occurred:\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-3-3b8a479202a4>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mdrive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmount\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/content/drive'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/google/colab/drive.py\u001b[0m in \u001b[0;36mmount\u001b[0;34m(mountpoint, force_remount, timeout_ms, use_metadata_server)\u001b[0m\n\u001b[1;32m    260\u001b[0m       \u001b[0;32mwith\u001b[0m \u001b[0m_output\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muse_tags\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'dfs-auth-dance'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfifo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'w'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mfifo_file\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m           \u001b[0mfifo_file\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwrite\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mget_code\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mauth_prompt\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'\\n'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\u001b[0m       \u001b[0mwrote_to_fifo\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mcase\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mraw_input\u001b[0;34m(self, prompt)\u001b[0m\n\u001b[1;32m    702\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_ident\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    703\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_header\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 704\u001b[0;31m             \u001b[0mpassword\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    705\u001b[0m         )\n\u001b[1;32m    706\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36m_input_request\u001b[0;34m(self, prompt, ident, parent, password)\u001b[0m\n\u001b[1;32m    732\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    733\u001b[0m                 \u001b[0;31m# re-raise KeyboardInterrupt, to truncate traceback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 734\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    735\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    736\u001b[0m                 \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TTj1RldNHgSQ"
      },
      "source": [
        "### Let's define some path, and our PokeMon dataset\n",
        "- Put the \"pokemon\" folder to somewhere of your Google Drive, and define the path to \"data_path\"\n",
        "- To 'model_dir', put the drive's directory path that you want to save your model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pcMuF_m9Ev8W"
      },
      "source": [
        "data_path = './drive/MyDrive/Dataset/pokemon' #./drive/MyDrive/Path/To/PokeMon/Which/Contains/Train/And/Validate\n",
        "model_dir = './drive/MyDrive/Codes/models'    #./drive/MyDrive/Path/To/Save/Your/Model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NFxwfpOxHcAk"
      },
      "source": [
        "class PokemonDataset(Dataset):\n",
        "    def __init__(self, data_path, is_training):\n",
        "        self.data_path = data_path\n",
        "        self.train_path = os.path.join(data_path, 'train')\n",
        "        self.val_path = os.path.join(data_path, 'validate')\n",
        "        self.is_training = is_training\n",
        "        if self.is_training:\n",
        "            self.target_path = self.train_path\n",
        "        else:\n",
        "            self.target_path = self.val_path\n",
        "\n",
        "        self.classes = sorted(os.listdir(self.target_path))\n",
        "        self.img_path_label = list()\n",
        "\n",
        "        for c in self.classes:\n",
        "            img_list = os.listdir(os.path.join(self.target_path, c))\n",
        "            for fp in img_list:\n",
        "                full_fp = os.path.join(self.target_path, c, fp)\n",
        "                self.img_path_label.append((full_fp, c, self.classes.index(c)))\n",
        "            \n",
        "        # Add some tranforms for data augmentation.\n",
        "        self.tensor_transform = torchvision.transforms.ToTensor()\n",
        "        self.normalize_transform = torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
        "                                                                    std=[0.229, 0.224, 0.225])\n",
        "        self.random_crop = torchvision.transforms.RandomCrop(size = 170)\n",
        "        self.random_flip = torchvision.transforms.RandomHorizontalFlip(p=0.5)\n",
        "        self.resize = torchvision.transforms.Resize(size=224)\n",
        "        self.train_transform = torchvision.transforms.Compose([self.tensor_transform,\n",
        "                                                            #    self.random_crop,\n",
        "                                                         self.random_flip,\n",
        "                                                         self.resize,\n",
        "                                                         self.normalize_transform])\n",
        "        self.validate_transform = torchvision.transforms.Compose([self.tensor_transform,\n",
        "                                                                  self.normalize_transform])\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.img_path_label)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        (fp, class_name, class_label) = self.img_path_label[idx]\n",
        "        img = Image.open(fp)\n",
        "        original_img = self.tensor_transform(img)\n",
        "\n",
        "        if self.is_training:\n",
        "            input = self.train_transform(img)\n",
        "        else:\n",
        "            input = self.validate_transform(img)\n",
        "            \n",
        "        sample = dict()\n",
        "        sample['input'] = input\n",
        "        sample['original_img'] = original_img\n",
        "        sample['target'] = class_label\n",
        "        sample['class_name'] = class_name\n",
        "\n",
        "        return sample"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UPFSRXfal8Dl"
      },
      "source": [
        "### Set DataSet and DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dJLrwTOxI127"
      },
      "source": [
        "batch_size = 64\n",
        "\n",
        "train_dataset = PokemonDataset(data_path, True)\n",
        "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)\n",
        "val_dataset = PokemonDataset(data_path, False)\n",
        "val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, pin_memory=True)\n",
        "\n",
        "num_classes = 18"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S6s35e4ymD-M"
      },
      "source": [
        "### Take a sample and try to look at the one"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pI2z7aqgMcUJ"
      },
      "source": [
        "sample = next(iter(train_dataloader))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6rsIBq2rIxA_"
      },
      "source": [
        "fig, ax = plt.subplots(1, 7, figsize=(20, 10))\n",
        "for i in range(7):\n",
        "    ax[i].imshow(sample['input'][i].permute(1, 2, 0))\n",
        "    ax[i].set_title(sample['class_name'][i])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PYiku-SxHnA6"
      },
      "source": [
        "### Choose your device - use GPU or not?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k44xmCOHHmF4"
      },
      "source": [
        "# device = 'cpu'\n",
        "device = 'cuda'\n",
        "print('Current Device : {}'.format(device))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LEjuX7lWIYlb"
      },
      "source": [
        "### Define the model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cx-fxCW5JD_w"
      },
      "source": [
        "class Model(nn.Module):\n",
        "    def __init__(self, feat_dim = 2048, dim_output=18):\n",
        "        super(Model, self).__init__()\n",
        "\n",
        "        self.dim_output = dim_output\n",
        "        self.feat_dim = feat_dim\n",
        "\n",
        "        self.conv1 = nn.Conv2d(feat_dim, feat_dim, kernel_size=1)\n",
        "        self.fc1 = nn.Linear(feat_dim, feat_dim//4) # 2048 -> 512\n",
        "        self.fc2 =  nn.Linear(feat_dim//4, feat_dim//8)\n",
        "        self.fc3 =  nn.Linear(feat_dim//8, dim_output)\n",
        "        self.relu = nn.LeakyReLU(0.1, inplace=True)\n",
        "        self.dropout1 = nn.Dropout(0.2)\n",
        "        self.dropout2 = nn.Dropout(0.5)\n",
        "        \n",
        "        \n",
        "        self.backbone = torch.hub.load('pytorch/vision:v0.6.0', 'resnext50_32x4d', pretrained=True)\n",
        "        # # Fix Initial Layers\n",
        "        for p in list(self.backbone.children())[:-5]:\n",
        "            p.requires_grad = False\n",
        "        # # get the structure until the last FC layer\n",
        "        modules = list(self.backbone.children())[:-1]\n",
        "        \n",
        "        self.backbone = nn.Sequential(*modules)\n",
        "    \n",
        "    def forward(self, img):\n",
        "        batch_size = img.shape[0]\n",
        "        x = self.backbone(img)\n",
        "        x = self.relu(self.conv1(x))\n",
        "        x = self.dropout1(x)\n",
        "        x = self.relu(self.fc1(x.view(batch_size, -1)))\n",
        "        x = self.dropout2(x)\n",
        "        x = self.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4hjFl4-IsqO"
      },
      "source": [
        "### Create a model and its optimizer\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZWtw-y-MKuv1"
      },
      "source": [
        "model = Model()\n",
        "model = model.to(device)\n",
        "\n",
        "optimizer = optim.AdamW(model.parameters(), lr=1e-4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JFVtNhFKuzLB"
      },
      "source": [
        "model(sample['input'].to(device)).shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DUSpLgVXIwqG"
      },
      "source": [
        "### Define functions for train/validation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q9pXkpAMK42y"
      },
      "source": [
        "def train(model, optimizer, sample):\n",
        "    model.train()\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    input = sample['input'].float().to(device)\n",
        "    target = sample['target'].long().to(device) \n",
        "    \n",
        "    pred = model(input)\n",
        "    pred_loss = criterion(pred, target)\n",
        "    \n",
        "    top3_val, top3_idx = torch.topk(pred, 3)\n",
        "\n",
        "    num_correct = torch.sum(top3_idx == target.view(-1, 1))\n",
        "    \n",
        "    pred_loss.backward()\n",
        "       \n",
        "    optimizer.step()\n",
        "\n",
        "    return pred_loss.item(), num_correct.item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MhcKwq8WOUGu"
      },
      "source": [
        "def validate(model, sample):\n",
        "    model.eval()\n",
        "\n",
        "    criterion = nn.CrossEntropyLoss()\n",
        "\n",
        "    with torch.no_grad():\n",
        "        input = sample['input'].float().to(device)\n",
        "        target = sample['target'].long().to(device) \n",
        "\n",
        "        pred = model(input)\n",
        "        pred_loss = criterion(pred, target)\n",
        "\n",
        "        top3_val, top3_idx = torch.topk(pred, 3)\n",
        "\n",
        "        num_correct = torch.sum(top3_idx == target.view(-1, 1))\n",
        "\n",
        "    return pred_loss.item(), num_correct.item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GjfAYAPsI2j4"
      },
      "source": [
        "### Prepare the Tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DLjrayaAMyb5"
      },
      "source": [
        "train_log_dir = './runs/train'\n",
        "train_summary_writer = summary.create_file_writer(train_log_dir)\n",
        "val_log_dir = './runs/validate'\n",
        "val_summary_writer = summary.create_file_writer(val_log_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rb5fy_VVRjUT"
      },
      "source": [
        "%tensorboard --logdir runs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T23Toz8vI7Gu"
      },
      "source": [
        "### Run Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-kTrpt4YMRS6"
      },
      "source": [
        "max_epoch = 200\n",
        "save_stride = 10\n",
        "tmp_path = './checkpoint.pth'\n",
        "max_accu = -1\n",
        "for epoch in tqdm(range(max_epoch)):        \n",
        "    ###Train Phase\n",
        "    \n",
        "    # Initialize Loss and Accuracy\n",
        "    train_loss = 0.0\n",
        "    train_accu = 0.0\n",
        "\n",
        "    # Load the saved MODEL AND OPTIMIZER after evaluation.\n",
        "    if epoch > 0:\n",
        "        checkpoint = torch.load(tmp_path)\n",
        "        model.load_state_dict(checkpoint['model_state_dict'])\n",
        "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
        "        # how about learning rate scheduler?\n",
        "\n",
        "    # Iterate over the train_dataloader\n",
        "    with tqdm(total=len(train_dataloader)) as pbar:\n",
        "        for idx, sample in enumerate(train_dataloader):\n",
        "            curr_loss, num_correct = train(model, optimizer, sample)\n",
        "            train_loss += curr_loss / len(train_dataloader)\n",
        "            train_accu += num_correct / len(train_dataset)\n",
        "            pbar.update(1)\n",
        "\n",
        "    # Write the current loss and accuracy to the Tensorboard\n",
        "    with train_summary_writer.as_default():\n",
        "        tf.summary.scalar('loss', train_loss, step=epoch)                \n",
        "        tf.summary.scalar('accuracy', train_accu, step=epoch)                \n",
        "\n",
        "    # save the model and optimizer's information before the evaulation\n",
        "    checkpoint = {\n",
        "        'model' : Model(),\n",
        "        'model_state_dict': model.state_dict(),\n",
        "        'optimizer_state_dict': optimizer.state_dict(),\n",
        "    }\n",
        "    \n",
        "    # Save the checkpoint - you can try to save the \"best\" model with the validation accuracy/loss\n",
        "    torch.save(checkpoint, tmp_path)\n",
        "    if (epoch+1) % save_stride == 0:\n",
        "        torch.save(checkpoint, os.path.join(model_dir, 'pokemon_{}.pth'.format(epoch+1)))\n",
        "    torch.save(checkpoint, os.path.join(model_dir, 'pokemon_recent.pth'))\n",
        "    \n",
        "    ### Validation Phase\n",
        "    # Initialize Loss and Accuracy\n",
        "    val_loss = 0.0\n",
        "    val_accu = 0.0\n",
        "\n",
        "    # Iterate over the val_dataloader\n",
        "    with tqdm(total=len(val_dataloader)) as pbar:\n",
        "        for idx, sample in enumerate(val_dataloader):\n",
        "            curr_loss, num_correct = validate(model, sample)\n",
        "            val_loss += curr_loss / len(val_dataloader)\n",
        "            val_accu += num_correct / len(val_dataloader)\n",
        "            pbar.update(1)\n",
        "\n",
        "    # Write the current loss and accuracy to the Tensorboard\n",
        "    with val_summary_writer.as_default():\n",
        "        tf.summary.scalar('loss', val_loss, step=epoch)\n",
        "        tf.summary.scalar('accuracy', val_accu, step=epoch) \n",
        "\n",
        "    max_accu = max(val_accu, max_accu)\n",
        "    if max_accu == val_accu:\n",
        "        # Save your best model to the checkpoint\n",
        "        torch.save(checkpoint, os.path.join(model_dir, 'pokemon_best.pth'))\n",
        "\n",
        "    # These Lines would make you update your Google Drive after the saving.\n",
        "    drive.flush_and_unmount()\n",
        "    drive.mount('/content/drive')\n",
        "\n",
        "    print(train_accu, val_accu)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5hPiA_pBdJk1"
      },
      "source": [
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dOeYopHTkhvJ"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.flush_and_unmount()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}