quickdraw.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "quickdraw.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMbD9KilU/oVAeLMDVOC9Ye",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/nateraw/quickdraw-pytorch/blob/main/quickdraw.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WMw65hcGEvr9"
      },
      "source": [
        "%%capture\n",
        "! pip install transformers wandb"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-2s4a4FU9P27"
      },
      "source": [
        "! wandb login"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ABuIopvX7B7V"
      },
      "source": [
        "import wandb"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LLce5mZO7FcZ"
      },
      "source": [
        "wandb.login()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qaviYKSbE3Pu"
      },
      "source": [
        "from typing import List, Optional\n",
        "import urllib.request\n",
        "from tqdm.auto import tqdm\n",
        "from pathlib import Path\n",
        "import requests\n",
        "import torch\n",
        "import math\n",
        "import numpy as np\n",
        "import os\n",
        "import glob\n",
        "\n",
        "\n",
        "def get_quickdraw_class_names():\n",
        "    \"\"\"\n",
        "    TODO - Check performance w/ gsutil in colab. The following command downloads all files to ./data\n",
        "    `gsutil cp gs://quickdraw_dataset/full/numpy_bitmap/* ./data`\n",
        "    \"\"\"\n",
        "    url = \"https://raw.githubusercontent.com/googlecreativelab/quickdraw-dataset/master/categories.txt\"\n",
        "    r = requests.get(url)\n",
        "    classes = [x.replace(' ', '_') for x in r.text.splitlines()]\n",
        "    return classes\n",
        "\n",
        "\n",
        "def download_quickdraw_dataset(root=\"./data\", limit: Optional[int] = None, class_names: List[str]=None):\n",
        "    if class_names is None:\n",
        "        class_names = get_quickdraw_class_names()\n",
        "\n",
        "    root = Path(root)\n",
        "    root.mkdir(exist_ok=True, parents=True)\n",
        "    url = 'https://storage.googleapis.com/quickdraw_dataset/full/numpy_bitmap/'\n",
        "\n",
        "    print(\"Downloading Quickdraw Dataset...\")\n",
        "    for class_name in tqdm(class_names[:limit]):\n",
        "        fpath = root / f\"{class_name}.npy\"\n",
        "        if not fpath.exists():\n",
        "            urllib.request.urlretrieve(f\"{url}{class_name.replace('_', '%20')}.npy\", fpath)\n",
        "\n",
        "\n",
        "def load_quickdraw_data(root=\"./data\", max_items_per_class=5000):\n",
        "    all_files = Path(root).glob('*.npy')\n",
        "\n",
        "    x = np.empty([0, 784], dtype=np.uint8)\n",
        "    y = np.empty([0], dtype=np.long)\n",
        "    class_names = []\n",
        "\n",
        "    print(f\"Loading {max_items_per_class} examples for each class from the Quickdraw Dataset...\")\n",
        "    for idx, file in enumerate(tqdm(sorted(all_files))):\n",
        "        data = np.load(file, mmap_mode='r')\n",
        "        data = data[0: max_items_per_class, :]\n",
        "        labels = np.full(data.shape[0], idx)\n",
        "        x = np.concatenate((x, data), axis=0)\n",
        "        y = np.append(y, labels)\n",
        "\n",
        "        class_names.append(file.stem)\n",
        "\n",
        "    return x, y, class_names\n",
        "\n",
        "\n",
        "class QuickDrawDataset(torch.utils.data.Dataset):\n",
        "    def __init__(self, root, max_items_per_class=5000, class_limit=None):\n",
        "        super().__init__()\n",
        "        self.root = root\n",
        "        self.max_items_per_class = max_items_per_class\n",
        "        self.class_limit = class_limit\n",
        "\n",
        "        download_quickdraw_dataset(self.root, self.class_limit)\n",
        "        self.X, self.Y, self.classes = load_quickdraw_data(self.root, self.max_items_per_class)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        x = (self.X[idx] / 255.).astype(np.float32).reshape(1, 28, 28)\n",
        "        y = self.Y[idx]\n",
        "\n",
        "        return torch.from_numpy(x), y.item()\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.X)\n",
        "\n",
        "    def collate_fn(self, batch):\n",
        "        x = torch.stack([item[0] for item in batch])\n",
        "        y = torch.LongTensor([item[1] for item in batch])\n",
        "        return {'pixel_values': x, 'labels': y}\n",
        "    \n",
        "    def split(self, pct=0.1):\n",
        "        num_classes = len(self.classes)\n",
        "        indices = torch.randperm(len(self)).tolist()\n",
        "        n_val = math.floor(len(indices) * pct)\n",
        "        train_ds = torch.utils.data.Subset(self, indices[:-n_val])\n",
        "        val_ds = torch.utils.data.Subset(self, indices[-n_val:])\n",
        "        return train_ds, val_ds"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "80jqbPi2E-CL"
      },
      "source": [
        "import torch\n",
        "from transformers import Trainer\n",
        "from transformers.modeling_utils import ModelOutput\n",
        "\n",
        "\n",
        "class QuickDrawTrainer(Trainer):\n",
        "\n",
        "    def compute_loss(self, model, inputs, return_outputs=False):\n",
        "        logits = model(inputs[\"pixel_values\"])\n",
        "        labels = inputs.get(\"labels\")\n",
        "\n",
        "        loss = None\n",
        "        if labels is not None:\n",
        "            loss_fct = torch.nn.CrossEntropyLoss()\n",
        "            loss = loss_fct(logits, labels)\n",
        "\n",
        "        return (loss, ModelOutput(logits=logits, loss=loss)) if return_outputs else loss\n",
        "\n",
        "# Taken from timm - https://github.com/rwightman/pytorch-image-models/blob/master/timm/utils/metrics.py\n",
        "def accuracy(output, target, topk=(1,)):\n",
        "    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n",
        "    maxk = min(max(topk), output.size()[1])\n",
        "    batch_size = target.size(0)\n",
        "    _, pred = output.topk(maxk, 1, True, True)\n",
        "    pred = pred.t()\n",
        "    correct = pred.eq(target.reshape(1, -1).expand_as(pred))\n",
        "    return [correct[:min(k, maxk)].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]\n",
        "\n",
        "\n",
        "def quickdraw_compute_metrics(p):\n",
        "    acc1, acc5 = accuracy(\n",
        "        torch.tensor(p.predictions),\n",
        "        torch.tensor(p.label_ids), topk=(1, 5)\n",
        "    )\n",
        "    return {'acc1': acc1, 'acc5': acc5}"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 115
        },
        "id": "L_tgy3PyFBpN",
        "outputId": "4ffc322c-2388-4945-e960-dfdcf2392318"
      },
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from transformers import TrainingArguments\n",
        "from datetime import datetime\n",
        "\n",
        "data_dir = './data'\n",
        "max_examples_per_class = 20000\n",
        "train_val_split_pct = .1\n",
        "\n",
        "ds = QuickDrawDataset(data_dir, max_examples_per_class)\n",
        "num_classes = len(ds.classes)\n",
        "train_ds, val_ds = ds.split(train_val_split_pct)"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading Quickdraw Dataset...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1fa4b4cfb6604655870a7d3e9c36f9b8",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "  0%|          | 0/345 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading 20000 examples for each class from the Quickdraw Dataset...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1644ed9448094038b9eb3b1fdaaa05bc",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "  0%|          | 0/345 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UpRtWkeH_lG5"
      },
      "source": [
        "## Define the Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IDuedQdSJF9z"
      },
      "source": [
        "# Original model used for hf.co/spaces/nateraw/quickdraw\n",
        "# When I trained the model for the spaces demo (link above), I limited to 100 classes. Full dataset is >300 classes\n",
        "# I'm updating to use the model uncommented below, as its slightly bigger and works better on the full dataset. \n",
        "# model = nn.Sequential(\n",
        "#     nn.Conv2d(1, 32, 3, padding='same'),\n",
        "#     nn.ReLU(),\n",
        "#     nn.MaxPool2d(2),\n",
        "#     nn.Conv2d(32, 64, 3, padding='same'),\n",
        "#     nn.ReLU(),\n",
        "#     nn.MaxPool2d(2),\n",
        "#     nn.Conv2d(64, 128, 3, padding='same'),\n",
        "#     nn.ReLU(),\n",
        "#     nn.MaxPool2d(2),\n",
        "#     nn.Flatten(),\n",
        "#     nn.Linear(1152, 512),\n",
        "#     nn.ReLU(),\n",
        "#     nn.Linear(512, num_classes),  # num_classes was limited to 100 here\n",
        "# )\n",
        "\n",
        "model = nn.Sequential(\n",
        "    nn.Conv2d(1, 64, 3, padding='same'),\n",
        "    nn.ReLU(),\n",
        "    nn.MaxPool2d(2),\n",
        "    nn.Conv2d(64, 128, 3, padding='same'),\n",
        "    nn.ReLU(),\n",
        "    nn.MaxPool2d(2),\n",
        "    nn.Conv2d(128, 256, 3, padding='same'),\n",
        "    nn.ReLU(),\n",
        "    nn.MaxPool2d(2),\n",
        "    nn.Flatten(),\n",
        "    nn.Linear(2304, 512),\n",
        "    nn.ReLU(),\n",
        "    nn.Linear(512, num_classes),\n",
        ")"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N4pzdmBK_mvt"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 696
        },
        "id": "7A9iOabpFJ6u",
        "outputId": "2f69c6d0-f45c-48eb-fed3-688992f46b04"
      },
      "source": [
        "timestamp = datetime.now().strftime('%Y-%m-%d-%H%M%S')\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=f'./outputs_20k_{timestamp}',\n",
        "    evaluation_strategy='epoch',\n",
        "    save_strategy='epoch',\n",
        "    report_to=['wandb', 'tensorboard'],  # Update to just tensorboard if not using wandb\n",
        "    logging_strategy='steps',\n",
        "    logging_steps=100,\n",
        "    per_device_train_batch_size=256,\n",
        "    per_device_eval_batch_size=256,\n",
        "    learning_rate=0.003,\n",
        "    fp16=torch.cuda.is_available(),\n",
        "    num_train_epochs=20,\n",
        "    run_name=f\"quickdraw-med-{timestamp}\",  # Can remove if not using wandb\n",
        "    warmup_steps=10000,\n",
        "    save_total_limit=5,\n",
        ")\n",
        "\n",
        "trainer = QuickDrawTrainer(\n",
        "    model,\n",
        "    training_args,\n",
        "    data_collator=ds.collate_fn,\n",
        "    train_dataset=train_ds,\n",
        "    eval_dataset=val_ds,\n",
        "    tokenizer=None,\n",
        "    compute_metrics=quickdraw_compute_metrics,\n",
        ")\n",
        "\n",
        "try:\n",
        "    # Training\n",
        "    train_results = trainer.train()\n",
        "    trainer.save_model()\n",
        "    trainer.log_metrics(\"train\", train_results.metrics)\n",
        "    trainer.save_metrics(\"train\", train_results.metrics)\n",
        "    trainer.save_state()\n",
        "\n",
        "    # Evaluation\n",
        "    eval_results = trainer.evaluate()\n",
        "    trainer.log_metrics(\"eval\", eval_results)\n",
        "    trainer.save_metrics(\"eval\", eval_results)\n",
        "except:\n",
        "    pass\n",
        "finally:\n",
        "    wandb.finish()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "PyTorch: setting up devices\n",
            "Using amp fp16 backend\n",
            "***** Running training *****\n",
            "  Num examples = 6210000\n",
            "  Num Epochs = 20\n",
            "  Instantaneous batch size per device = 256\n",
            "  Total train batch size (w. parallel, distributed & accumulation) = 256\n",
            "  Gradient Accumulation steps = 1\n",
            "  Total optimization steps = 485160\n",
            "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                    Syncing run <strong><a href=\"https://wandb.ai/nateraw/huggingface/runs/2pha0ydi\" target=\"_blank\">quickdraw-med-2021-10-12-210155</a></strong> to <a href=\"https://wandb.ai/nateraw/huggingface\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">docs</a>).<br/>\n",
              "\n",
              "                "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='72775' max='485160' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [ 72770/485160 19:18 < 1:49:27, 62.79 it/s, Epoch 3.00/20]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: left;\">\n",
              "      <th>Epoch</th>\n",
              "      <th>Training Loss</th>\n",
              "      <th>Validation Loss</th>\n",
              "      <th>Acc1</th>\n",
              "      <th>Acc5</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>1</td>\n",
              "      <td>1.429900</td>\n",
              "      <td>1.424564</td>\n",
              "      <td>65.384491</td>\n",
              "      <td>87.336525</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2</td>\n",
              "      <td>1.379100</td>\n",
              "      <td>1.399322</td>\n",
              "      <td>65.749565</td>\n",
              "      <td>87.547104</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>\n",
              "    <div>\n",
              "      \n",
              "      <progress value='92' max='2696' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [  92/2696 00:00 < 00:20, 125.68 it/s]\n",
              "    </div>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py:1357: FutureWarning: Non-finite norm encountered in torch.nn.utils.clip_grad_norm_; continuing anyway. Note that the default behavior will change in a future release to error out if a non-finite total norm is encountered. At that point, setting error_if_nonfinite=false will be required to retain the old behavior.\n",
            "  args.max_grad_norm,\n",
            "***** Running Evaluation *****\n",
            "  Num examples = 690000\n",
            "  Batch size = 256\n",
            "Saving model checkpoint to ./outputs_20k_2021-10-12-210155/checkpoint-24258\n",
            "Trainer.model is not a `PreTrainedModel`, only saving its state dict.\n",
            "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py:1357: FutureWarning: Non-finite norm encountered in torch.nn.utils.clip_grad_norm_; continuing anyway. Note that the default behavior will change in a future release to error out if a non-finite total norm is encountered. At that point, setting error_if_nonfinite=false will be required to retain the old behavior.\n",
            "  args.max_grad_norm,\n",
            "***** Running Evaluation *****\n",
            "  Num examples = 690000\n",
            "  Batch size = 256\n",
            "Saving model checkpoint to ./outputs_20k_2021-10-12-210155/checkpoint-48516\n",
            "Trainer.model is not a `PreTrainedModel`, only saving its state dict.\n",
            "/usr/local/lib/python3.7/dist-packages/transformers/trainer.py:1357: FutureWarning: Non-finite norm encountered in torch.nn.utils.clip_grad_norm_; continuing anyway. Note that the default behavior will change in a future release to error out if a non-finite total norm is encountered. At that point, setting error_if_nonfinite=false will be required to retain the old behavior.\n",
            "  args.max_grad_norm,\n",
            "***** Running Evaluation *****\n",
            "  Num examples = 690000\n",
            "  Batch size = 256\n"
          ]
        }
      ]
    }
  ]
}