Pytorch/Prototypical Networks for Few-Shot-Learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation of Prototypical Networks for Few Shot Learning (https://arxiv.org/abs/1703.05175) in Pytorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Prototypical Networks\n",
    "> ImageNet’s dataset contains a multitude of examples for machine learning, which is not always the case\n",
    "> And typical deep learning architecture relies on substantial data for sufficient outcomes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> With N-shot learning, we want to create our model which is trained using a very limited number of images (less than N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Prototypical networks are one of the most popular deep learning algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Unlike typical deep learning architecture, prototypical networks do not classify the image directly, and instead learn the mapping of an image in metric space.\n",
    "> It simply calculates the distance betwwen a vector in an image (we only compute the distance of one point to another), so the \"origin\" point does not matter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['images_background', 'images_evaluation']\n"
     ]
    }
   ],
   "source": [
    "import multiprocessing as mp\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import numpy as np \n",
    "import pandas as pd\n",
    "import os\n",
    "print(os.listdir(\"C:/Users/bokhy/Desktop/Python/Python-Projects/Pytorch/data/Omniglot\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Read / Import Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The network was trained on the Omniglot dataset. The Omniglot data set is designed for developing more human-like learning algorithms. It contains 1,623 different handwritten characters from 50 different alphabets\n",
    "> Data: https://github.com/brendenlake/omniglot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dir = os.listdir('C:/Users/bokhy/Desktop/Python/Python-Projects/Pytorch/data/Omniglot/images_background/')\n",
    "datax = np.array([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Alphabet_of_the_Magi',\n",
       " 'Anglo-Saxon_Futhorc',\n",
       " 'Arcadian',\n",
       " 'Armenian',\n",
       " 'Asomtavruli_(Georgian)',\n",
       " 'Balinese',\n",
       " 'Bengali',\n",
       " 'Blackfoot_(Canadian_Aboriginal_Syllabics)',\n",
       " 'Braille',\n",
       " 'Burmese_(Myanmar)',\n",
       " 'Cyrillic',\n",
       " 'Early_Aramaic',\n",
       " 'Futurama',\n",
       " 'Grantha',\n",
       " 'Greek',\n",
       " 'Gujarati',\n",
       " 'Hebrew',\n",
       " 'Inuktitut_(Canadian_Aboriginal_Syllabics)',\n",
       " 'Japanese_(hiragana)',\n",
       " 'Japanese_(katakana)',\n",
       " 'Korean',\n",
       " 'Latin',\n",
       " 'Malay_(Jawi_-_Arabic)',\n",
       " 'Mkhedruli_(Georgian)',\n",
       " 'N_Ko',\n",
       " 'Ojibwe_(Canadian_Aboriginal_Syllabics)',\n",
       " 'Sanskrit',\n",
       " 'Syriac_(Estrangelo)',\n",
       " 'Tagalog',\n",
       " 'Tifinagh']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Data Augmentation: Images are rotated by 90, 180 and 270 degrees and each rotation resulted in one more class. Hence the total count of classes reached to 6492 (1,623 * 4) total classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> TODO: Use Albumentation library to augment images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_rotate(img, angle):\n",
    "    rows,cols, _ = img.shape\n",
    "    M = cv2.getRotationMatrix2D((cols/2 ,rows/2),angle,1)\n",
    "    dst = cv2.warpAffine(img,M,(cols,rows))\n",
    "    return np.expand_dims(dst, 0)\n",
    "\n",
    "def read_alphabets(alphabet_directory, directory):\n",
    "    \"\"\"\n",
    "    Reads all the characters from alphabet_directory and augment each image with 90, 180, 270 degrees of rotation.\n",
    "    \"\"\"\n",
    "    datax = None\n",
    "    datay = []\n",
    "    characters = os.listdir(alphabet_directory)\n",
    "    for character in characters:\n",
    "        images = os.listdir(alphabet_directory + character + '/')\n",
    "        for img in images:\n",
    "            image = cv2.resize(cv2.imread(alphabet_directory + character + '/' + img), (28,28))\n",
    "            image90 = image_rotate(image, 90)\n",
    "            image180 = image_rotate(image, 180)\n",
    "            image270 = image_rotate(image, 270)\n",
    "            image = np.expand_dims(image, 0)\n",
    "            if datax is None:\n",
    "                datax = np.vstack([image, image90, image180, image270])\n",
    "            else:\n",
    "                datax = np.vstack([datax, image, image90, image180, image270])\n",
    "            datay.append(directory + '_' + character + '_0')\n",
    "            datay.append(directory + '_' + character + '_90')\n",
    "            datay.append(directory + '_' + character + '_180')\n",
    "            datay.append(directory + '_' + character + '_270')\n",
    "    return datax, np.array(datay)\n",
    "\n",
    "def read_images(base_directory):\n",
    "    \"\"\"\n",
    "    Used multithreading for data reading to decrease the reading time drastically\n",
    "    \"\"\"\n",
    "    datax = None\n",
    "    datay = []\n",
    "    pool = mp.Pool(mp.cpu_count())\n",
    "    results = [pool.apply(read_alphabets, args=(base_directory + '/' + directory + '/', directory, )) for directory in os.listdir(base_directory)]\n",
    "    pool.close()\n",
    "    for result in results:\n",
    "        if datax is None:\n",
    "            datax = result[0]\n",
    "            datay = result[1]\n",
    "        else:\n",
    "            datax = np.vstack([datax, result[0]])\n",
    "            datay = np.concatenate([datay, result[1]])\n",
    "    return datax, datay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 0 ns\n"
     ]
    }
   ],
   "source": [
    "%time trainx, trainy = read_images('C:/Users/bokhy/Desktop/Python/Python-Projects/Pytorch/data/Omniglot/images_background/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time testx, testy = read_images('C:/Users/bokhy/Desktop/Python/Python-Projects/Pytorch/data/Omniglot/images_evaluation/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Making sure we have shpaes that we want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainx.shape, trainy.shape, testx.shape, testy.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Create a model,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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.autograd import Variable\n",
    "from tqdm import trange\n",
    "from time import sleep\n",
    "from sklearn.preprocessing import OneHotEncoder, LabelEncoder\n",
    "use_gpu = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainx = torch.from_numpy(trainx).float()\n",
    "#trainy = torch.from_numpy(trainy)\n",
    "testx = torch.from_numpy(testx).float()\n",
    "#testy = torch.from_numpy(testy)\n",
    "if use_gpu:\n",
    "    trainx = trainx.to(device)\n",
    "    testx = testx.to(device)\n",
    "trainx.size(), testx.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### \"Image2vector CNN\" architecture. It takes images of dimensions 28x28x3 and returns a vector of length 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    \n",
    "    def sub_block(self, in_channels, out_channels=64, kernel_size=3):\n",
    "        block = torch.nn.Sequential(\n",
    "                    torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=out_channels, padding=1),\n",
    "                    torch.nn.BatchNorm2d(out_channels),\n",
    "                    torch.nn.ReLU(),\n",
    "                    torch.nn.MaxPool2d(kernel_size=2)\n",
    "                )\n",
    "        return block\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.convnet1 = self.sub_block(3)\n",
    "        self.convnet2 = self.sub_block(64)\n",
    "        self.convnet3 = self.sub_block(64)\n",
    "        self.convnet4 = self.sub_block(64)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.convnet1(x)\n",
    "        x = self.convnet2(x)\n",
    "        x = self.convnet3(x)\n",
    "        x = self.convnet4(x)\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrototypicalNet(nn.Module):\n",
    "    def __init__(self, use_gpu=False):\n",
    "        super(PrototypicalNet, self).__init__()\n",
    "        self.f = Net()\n",
    "        self.gpu = use_gpu\n",
    "        if self.gpu:\n",
    "            self.f = self.f.cuda()\n",
    "    \n",
    "    def forward(self, datax, datay, Ns,Nc, Nq, total_classes):\n",
    "        \"\"\"\n",
    "        Implementation of one episode in Prototypical Net\n",
    "        datax: Training images\n",
    "        datay: Corresponding labels of datax\n",
    "        Nc: Number  of classes per episode\n",
    "        Ns: Number of support data per class\n",
    "        Nq:  Number of query data per class\n",
    "        total_classes: Total classes in training set\n",
    "        \"\"\"\n",
    "        k = total_classes.shape[0]\n",
    "        K = np.random.choice(total_classes, Nc, replace=False)\n",
    "        Query_x = torch.Tensor()\n",
    "        if(self.gpu):\n",
    "            Query_x = Query_x.cuda()\n",
    "        Query_y = []\n",
    "        Query_y_count = []\n",
    "        centroid_per_class  = {}\n",
    "        class_label = {}\n",
    "        label_encoding = 0\n",
    "        for cls in K:\n",
    "            S_cls, Q_cls = self.random_sample_cls(datax, datay, Ns, Nq, cls)\n",
    "            centroid_per_class[cls] = self.get_centroid(S_cls, Nc)\n",
    "            class_label[cls] = label_encoding\n",
    "            label_encoding += 1\n",
    "            Query_x = torch.cat((Query_x, Q_cls), 0) # Joining all the query set together\n",
    "            Query_y += [cls]\n",
    "            Query_y_count += [Q_cls.shape[0]]\n",
    "        Query_y, Query_y_labels = self.get_query_y(Query_y, Query_y_count, class_label)\n",
    "        Query_x = self.get_query_x(Query_x, centroid_per_class, Query_y_labels)\n",
    "        return Query_x, Query_y\n",
    "    \n",
    "    def random_sample_cls(self, datax, datay, Ns, Nq, cls):\n",
    "        \"\"\"\n",
    "        Randomly samples Ns examples as support set and Nq as Query set\n",
    "        \"\"\"\n",
    "        data = datax[(datay == cls).nonzero()]\n",
    "        perm = torch.randperm(data.shape[0])\n",
    "        idx = perm[:Ns]\n",
    "        S_cls = data[idx]\n",
    "        idx = perm[Ns : Ns+Nq]\n",
    "        Q_cls = data[idx]\n",
    "        if self.gpu:\n",
    "            S_cls = S_cls.cuda()\n",
    "            Q_cls = Q_cls.cuda()\n",
    "        return S_cls, Q_cls\n",
    "    \n",
    "    def get_centroid(self, S_cls, Nc):\n",
    "        \"\"\"\n",
    "        Returns a centroid vector of support set for a class\n",
    "        \"\"\"\n",
    "        return torch.sum(self.f(S_cls), 0).unsqueeze(1).transpose(0,1) / Nc\n",
    "    \n",
    "    def get_query_y(self, Qy, Qyc, class_label):\n",
    "        \"\"\"\n",
    "        Returns labeled representation of classes of Query set and a list of labels.\n",
    "        \"\"\"\n",
    "        labels = []\n",
    "        m = len(Qy)\n",
    "        for i in range(m):\n",
    "            labels += [Qy[i]] * Qyc[i]\n",
    "        labels = np.array(labels).reshape(len(labels), 1)\n",
    "        label_encoder = LabelEncoder()\n",
    "        Query_y = torch.Tensor(label_encoder.fit_transform(labels).astype(int)).long()\n",
    "        if self.gpu:\n",
    "            Query_y = Query_y.cuda()\n",
    "        Query_y_labels = np.unique(labels)\n",
    "        return Query_y, Query_y_labels\n",
    "    \n",
    "    def get_centroid_matrix(self, centroid_per_class, Query_y_labels):\n",
    "        \"\"\"\n",
    "        Returns the centroid matrix where each column is a centroid of a class.\n",
    "        \"\"\"\n",
    "        centroid_matrix = torch.Tensor()\n",
    "        if(self.gpu):\n",
    "            centroid_matrix = centroid_matrix.cuda()\n",
    "        for label in Query_y_labels:\n",
    "            centroid_matrix = torch.cat((centroid_matrix, centroid_per_class[label]))\n",
    "        if self.gpu:\n",
    "            centroid_matrix = centroid_matrix.cuda()\n",
    "        return centroid_matrix\n",
    "    \n",
    "    def get_query_x(self, Query_x, centroid_per_class, Query_y_labels):\n",
    "        \"\"\"\n",
    "        Returns distance matrix from each Query image to each centroid.\n",
    "        \"\"\"\n",
    "        centroid_matrix = self.get_centroid_matrix(centroid_per_class, Query_y_labels)\n",
    "        Query_x = self.f(Query_x)\n",
    "        m = Query_x.size(0)\n",
    "        n = centroid_matrix.size(0)\n",
    "        # The below expressions expand both the matrices such that they become compatible to each other in order to caclulate L2 distance.\n",
    "        centroid_matrix = centroid_matrix.expand(m, centroid_matrix.size(0), centroid_matrix.size(1)) # Expanding centroid matrix to \"m\".\n",
    "        Query_matrix = Query_x.expand(n, Query_x.size(0), Query_x.size(1)).transpose(0,1) # Expanding Query matrix \"n\" times\n",
    "        Qx = torch.pairwise_distance(centroid_matrix.transpose(1,2), Query_matrix.transpose(1,2))\n",
    "        return Qx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "protonet = PrototypicalNet(use_gpu=use_gpu)\n",
    "optimizer = optim.SGD(protonet.parameters(), lr = 0.01, momentum=0.99)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_step(datax, datay, Ns,Nc, Nq):\n",
    "    optimizer.zero_grad()\n",
    "    Qx, Qy= protonet(datax, datay, Ns, Nc, Nq, np.unique(datay))\n",
    "    pred = torch.log_softmax(Qx, dim=-1)\n",
    "    loss = F.nll_loss(pred, Qy)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    acc = torch.mean((torch.argmax(pred, 1) == Qy).float())\n",
    "    return loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_episode = 16000\n",
    "frame_size = 1000\n",
    "frame_loss = 0\n",
    "frame_acc = 0\n",
    "\n",
    "for i in range(num_episode):\n",
    "    loss, acc = train_step(trainx, trainy, 5, 60, 5)\n",
    "    frame_loss += loss.data\n",
    "    frame_acc += acc.data\n",
    "    if( (i+1) % frame_size == 0):\n",
    "        print(\"Frame Number:\", ((i+1) // frame_size), 'Frame Loss: ', frame_loss.data.cpu().numpy().tolist()/ frame_size, 'Frame Accuracy:', (frame_acc.data.cpu().numpy().tolist() * 100) / frame_size)\n",
    "        frame_loss = 0\n",
    "        frame_acc = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Test the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_step(datax, datay, Ns,Nc, Nq):\n",
    "    Qx, Qy= protonet(datax, datay, Ns, Nc, Nq, np.unique(datay))\n",
    "    pred = torch.log_softmax(Qx, dim=-1)\n",
    "    loss = F.nll_loss(pred, Qy)\n",
    "    acc = torch.mean((torch.argmax(pred, 1) == Qy).float())\n",
    "    return loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_test_episode = 2000\n",
    "avg_loss = 0\n",
    "avg_acc = 0\n",
    "\n",
    "for _ in range(num_test_episode):\n",
    "    loss, acc = test_step(testx, testy, 5, 60, 15)\n",
    "    avg_loss += loss.data\n",
    "    avg_acc += acc.data\n",
    "print('Avg Loss: ', avg_loss.data.cpu().numpy().tolist() / num_test_episode , 'Avg Accuracy:', (avg_acc.data.cpu().numpy().tolist() * 100) / num_test_episode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "> Weight-loading for later use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_weights(filename, protonet, use_gpu):\n",
    "    if use_gpu:\n",
    "        protonet.load_state_dict(torch.load(filename))\n",
    "    else:\n",
    "        protonet.load_state_dict(torch.load(filename), map_location='cpu')\n",
    "    return protonet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Limitations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> It works well on predicting images of a character using the same dataset, because they share the similar traits. However, if we were to try using this model to use on different tpyes of image (i.e. fruits), it would not be as good as this one, as fruits and writings share few characteristics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Reference: https://blog.floydhub.com/n-shot-learning/"
   ]
  }
 ],
 "metadata": {
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}