class-norm-for-czsl.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Class_Normaliation_for_Continual_Zero_Shot_Learning.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.7.4"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/universome/class-norm-for-czsl/blob/master/class-norm-for-czsl.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ML0uxd0ft-ZV"
      },
      "source": [
        "#### 1. Defining the hyperparams"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t_Gl6w5ft-ZV"
      },
      "source": [
        "DATASET = 'AWA1' # One of [\"AWA1\", \"AWA2\", \"APY\", \"CUB\", \"SUN\"]\n",
        "USE_CLASS_STANDARTIZATION = True # i.e. equation (9) from the paper\n",
        "USE_PROPER_INIT = True # i.e. equation (10) from the paper"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nLyeh2lit-ZW"
      },
      "source": [
        "#### 2. Downloading GBU data from [the official GBU website](https://www.mpi-inf.mpg.de/departments/computer-vision-and-machine-learning/research/zero-shot-learning/zero-shot-learning-the-good-the-bad-and-the-ugly) (takes 1-2 minutes for the first time)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6qt2FThst-ZW"
      },
      "source": [
        "%%bash\n",
        "if [ -d \"./data\" ] \n",
        "then\n",
        "    echo \"Files are already there.\"\n",
        "else\n",
        "    wget -q \"http://datasets.d2.mpi-inf.mpg.de/xian/xlsa17.zip\"\n",
        "    unzip -q xlsa17.zip -d ./data\n",
        "fi"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XW6mEZNyt-ZX"
      },
      "source": [
        "#### 3. Running the code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jh0MKpUtt-ZX",
        "outputId": "9761732b-c765-43af-8fca-1ac71057f8c9"
      },
      "source": [
        "import numpy as np; np.random.seed(1)\n",
        "import torch; torch.manual_seed(1)\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from time import time\n",
        "from tqdm import tqdm\n",
        "from scipy import io\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "\n",
        "print(f'<=============== Loading data for {DATASET} ===============>')\n",
        "DEVICE = 'cuda' # Set to 'cpu' if a GPU is not available\n",
        "DATA_DIR = f'./data/xlsa17/data/{DATASET}'\n",
        "data = io.loadmat(f'{DATA_DIR}/res101.mat')\n",
        "attrs_mat = io.loadmat(f'{DATA_DIR}/att_splits.mat')\n",
        "feats = data['features'].T.astype(np.float32)\n",
        "labels = data['labels'].squeeze() - 1 # Using \"-1\" here and for idx to normalize to 0-index\n",
        "train_idx = attrs_mat['trainval_loc'].squeeze() - 1\n",
        "test_seen_idx = attrs_mat['test_seen_loc'].squeeze() - 1\n",
        "test_unseen_idx = attrs_mat['test_unseen_loc'].squeeze() - 1\n",
        "test_idx = np.array(test_seen_idx.tolist() + test_unseen_idx.tolist())\n",
        "seen_classes = sorted(np.unique(labels[test_seen_idx]))\n",
        "unseen_classes = sorted(np.unique(labels[test_unseen_idx]))\n",
        "\n",
        "\n",
        "print(f'<=============== Preprocessing ===============>')\n",
        "num_classes = len(seen_classes) + len(unseen_classes)\n",
        "seen_mask = np.array([(c in seen_classes) for c in range(num_classes)])\n",
        "unseen_mask = np.array([(c in unseen_classes) for c in range(num_classes)])\n",
        "attrs = attrs_mat['att'].T\n",
        "attrs = torch.from_numpy(attrs).to(DEVICE).float()\n",
        "attrs = attrs / attrs.norm(dim=1, keepdim=True) * np.sqrt(attrs.shape[1])\n",
        "attrs_seen = attrs[seen_mask]\n",
        "attrs_unseen = attrs[unseen_mask]\n",
        "train_labels = labels[train_idx]\n",
        "test_labels = labels[test_idx]\n",
        "test_seen_idx = [i for i, y in enumerate(test_labels) if y in seen_classes]\n",
        "test_unseen_idx = [i for i, y in enumerate(test_labels) if y in unseen_classes]\n",
        "labels_remapped_to_seen = [(seen_classes.index(t) if t in seen_classes else -1) for t in labels]\n",
        "test_labels_remapped_seen = [(seen_classes.index(t) if t in seen_classes else -1) for t in test_labels]\n",
        "test_labels_remapped_unseen = [(unseen_classes.index(t) if t in unseen_classes else -1) for t in test_labels]\n",
        "ds_train = [(feats[i], labels_remapped_to_seen[i]) for i in train_idx]\n",
        "ds_test = [(feats[i], int(labels[i])) for i in test_idx]\n",
        "train_dataloader = DataLoader(ds_train, batch_size=256, shuffle=True)\n",
        "test_dataloader = DataLoader(ds_test, batch_size=2048)\n",
        "\n",
        "class_indices_inside_test = {c: [i for i in range(len(test_idx)) if labels[test_idx[i]] == c] for c in range(num_classes)}\n",
        "\n",
        "\n",
        "class ClassStandardization(nn.Module):\n",
        "    \"\"\"\n",
        "    Class Standardization procedure from the paper.\n",
        "    Conceptually, it is equivalent to nn.BatchNorm1d with affine=False,\n",
        "    but for some reason nn.BatchNorm1d performs slightly worse.\n",
        "    \"\"\"\n",
        "    def __init__(self, feat_dim: int):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.running_mean = nn.Parameter(torch.zeros(feat_dim), requires_grad=False)\n",
        "        self.running_var = nn.Parameter(torch.ones(feat_dim), requires_grad=False)\n",
        "    \n",
        "    def forward(self, class_feats):\n",
        "        \"\"\"\n",
        "        Input: class_feats of shape [num_classes, feat_dim]\n",
        "        Output: class_feats (standardized) of shape [num_classes, feat_dim]\n",
        "        \"\"\"\n",
        "        if self.training:\n",
        "            batch_mean = class_feats.mean(dim=0)\n",
        "            batch_var = class_feats.var(dim=0)\n",
        "            \n",
        "            # Normalizing the batch\n",
        "            result = (class_feats - batch_mean.unsqueeze(0)) / (batch_var.unsqueeze(0) + 1e-5)\n",
        "            \n",
        "            # Updating the running mean/std\n",
        "            self.running_mean.data = 0.9 * self.running_mean.data + 0.1 * batch_mean.detach()\n",
        "            self.running_var.data = 0.9 * self.running_var.data + 0.1 * batch_var.detach()\n",
        "        else:\n",
        "            # Using accumulated statistics\n",
        "            # Attention! For the test inference, we cant use batch-wise statistics,\n",
        "            # only the accumulated ones. Otherwise, it will be quite transductive\n",
        "            result = (class_feats - self.running_mean.unsqueeze(0)) / (self.running_var.unsqueeze(0) + 1e-5)\n",
        "        \n",
        "        return result\n",
        "\n",
        "\n",
        "class CNZSLModel(nn.Module):\n",
        "    def __init__(self, attr_dim: int, hid_dim: int, proto_dim: int):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.model = nn.Sequential(\n",
        "            nn.Linear(attr_dim, hid_dim),\n",
        "            nn.ReLU(),\n",
        "            \n",
        "            nn.Linear(hid_dim, hid_dim),\n",
        "            ClassStandardization(hid_dim) if USE_CLASS_STANDARTIZATION else nn.Identity(),\n",
        "            nn.ReLU(),\n",
        "            \n",
        "            ClassStandardization(hid_dim) if USE_CLASS_STANDARTIZATION else nn.Identity(),\n",
        "            nn.Linear(hid_dim, proto_dim),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "        \n",
        "        if USE_PROPER_INIT:\n",
        "            weight_var = 1 / (hid_dim * proto_dim)\n",
        "            b = np.sqrt(3 * weight_var)\n",
        "            self.model[-2].weight.data.uniform_(-b, b)\n",
        "        \n",
        "    def forward(self, x, attrs):\n",
        "        protos = self.model(attrs)\n",
        "        x_ns = 5 * x / x.norm(dim=1, keepdim=True) # [batch_size, x_dim]\n",
        "        protos_ns = 5 * protos / protos.norm(dim=1, keepdim=True) # [num_classes, x_dim]\n",
        "        logits = x_ns @ protos_ns.t() # [batch_size, num_classes]\n",
        "        \n",
        "        return logits\n",
        "    \n",
        "\n",
        "print(f'\\n<=============== Starting training ===============>')\n",
        "start_time = time()\n",
        "model = CNZSLModel(attrs.shape[1], 1024, feats.shape[1]).to(DEVICE)\n",
        "optim = torch.optim.Adam(model.model.parameters(), lr=0.0005, weight_decay=0.0001)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optim, gamma=0.1, step_size=25)\n",
        "\n",
        "\n",
        "for epoch in tqdm(range(50)):\n",
        "    model.train()\n",
        "    \n",
        "    for i, batch in enumerate(train_dataloader):\n",
        "        feats = torch.from_numpy(np.array(batch[0])).to(DEVICE)\n",
        "        targets = torch.from_numpy(np.array(batch[1])).to(DEVICE)\n",
        "        logits = model(feats, attrs[seen_mask])\n",
        "        loss = F.cross_entropy(logits, targets)\n",
        "        optim.zero_grad()\n",
        "        loss.backward()\n",
        "        optim.step()\n",
        "    \n",
        "    scheduler.step()\n",
        "\n",
        "print(f'Training is done! Took time: {(time() - start_time): .1f} seconds')\n",
        "\n",
        "model.eval() # Important! Otherwise we would use unseen batch statistics\n",
        "logits = [model(x.to(DEVICE), attrs).cpu() for x, _ in test_dataloader]\n",
        "logits = torch.cat(logits, dim=0)\n",
        "logits[:, seen_mask] *= (0.95 if DATASET != \"CUB\" else 1.0) # Trading a bit of gzsl-s for a bit of gzsl-u\n",
        "preds_gzsl = logits.argmax(dim=1).numpy()\n",
        "preds_zsl_s = logits[:, seen_mask].argmax(dim=1).numpy()\n",
        "preds_zsl_u = logits[:, ~seen_mask].argmax(dim=1).numpy()\n",
        "guessed_zsl_u = (preds_zsl_u == test_labels_remapped_unseen)\n",
        "guessed_gzsl = (preds_gzsl == test_labels)\n",
        "zsl_unseen_acc = np.mean([guessed_zsl_u[cls_idx].mean().item() for cls_idx in [class_indices_inside_test[c] for c in unseen_classes]]) \n",
        "gzsl_seen_acc = np.mean([guessed_gzsl[cls_idx].mean().item() for cls_idx in [class_indices_inside_test[c] for c in seen_classes]])\n",
        "gzsl_unseen_acc = np.mean([guessed_gzsl[cls_idx].mean().item() for cls_idx in [class_indices_inside_test[c] for c in unseen_classes]])\n",
        "gzsl_harmonic = 2 * (gzsl_seen_acc * gzsl_unseen_acc) / (gzsl_seen_acc + gzsl_unseen_acc)\n",
        "\n",
        "print(f'ZSL-U: {zsl_unseen_acc * 100:.02f}')\n",
        "print(f'GZSL-U: {gzsl_unseen_acc * 100:.02f}')\n",
        "print(f'GZSL-S: {gzsl_seen_acc * 100:.02f}')\n",
        "print(f'GZSL-H: {gzsl_harmonic * 100:.02f}')"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "<=============== Loading data for AWA1 ===============>\n",
            "<=============== Preprocessing ===============>\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\r  0%|          | 0/50 [00:00<?, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "<=============== Starting training ===============>\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 50/50 [00:31<00:00,  1.61it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Training is done! Took time:  31.1 seconds\n",
            "ZSL-U: 69.70\n",
            "GZSL-U: 64.08\n",
            "GZSL-S: 73.41\n",
            "GZSL-H: 68.43\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stderr"
        }
      ]
    }
  ]
}