algorithm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import builtins\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import argparse\n",
    "import socket\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torchvision import transforms, datasets\n",
    "\n",
    "from PIL import ImageFilter\n",
    "\n",
    "import models.resnet as resnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_option():\n",
    "\n",
    "    parser = argparse.ArgumentParser('argument for training')\n",
    "\n",
    "    parser.add_argument('data', type=str, help='path to dataset')\n",
    "    parser.add_argument('--dataset', type=str, default='imagenet',\n",
    "                        choices=['imagenet', 'imagenet100'],\n",
    "                        help='use full or subset of the dataset')\n",
    "    parser.add_argument('--num_workers', type=int, default=12, help='num of workers to use')\n",
    "\n",
    "    # model definition\n",
    "    parser.add_argument('--arch', type=str, default='resnet18',\n",
    "                        choices=['resnet18' , 'resnet50', 'mobilenet'])\n",
    "\n",
    "    # ISD loss function\n",
    "    parser.add_argument('--queue_size', type=int, default=128000)\n",
    "    parser.add_argument('--temp', type=float, default=0.02)\n",
    "    parser.add_argument('--momentum', type=float, default=0.999)\n",
    "\n",
    "    parser.add_argument('--resume', default='', type=str,\n",
    "                        help='path to latest checkpoint (default: none)')\n",
    "    return parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extended version of ImageFolder to return index of image too.\n",
    "class ImageFolderEx(datasets.ImageFolder) :\n",
    "    def __getitem__(self, index):\n",
    "        sample, target = super(ImageFolderEx, self).__getitem__(index)\n",
    "        return index, sample, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform the images twice: one for query encoder and the other for key encoder\n",
    "class TwoCropsTransform:\n",
    "    \"\"\"Take two random crops of one image as the query and key.\"\"\"\n",
    "\n",
    "    def __init__(self, base_transform):\n",
    "        self.base_transform = base_transform\n",
    "\n",
    "    def __call__(self, x):\n",
    "        q = self.base_transform(x)\n",
    "        k = self.base_transform(x)\n",
    "        return [q, k]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the dataloader with two augmentations for each image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create train loader\n",
    "def get_train_loader(opt):\n",
    "    traindir = os.path.join(opt.data, 'train')\n",
    "    image_size = 224\n",
    "    mean = [0.485, 0.456, 0.406]\n",
    "    std = [0.229, 0.224, 0.225]\n",
    "    normalize = transforms.Normalize(mean=mean, std=std)\n",
    "\n",
    "    # use less severe augmentation than in the actual code\n",
    "    augmentation = [\n",
    "        transforms.RandomResizedCrop(224, scale=(0.2, 1.)),\n",
    "        transforms.RandomGrayscale(p=0.2),\n",
    "        transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        normalize\n",
    "    ]\n",
    "\n",
    "\n",
    "    train_dataset = ImageFolderEx(\n",
    "        traindir,\n",
    "        TwoCropsTransform(transforms.Compose(augmentation))\n",
    "    )\n",
    "    print('==> train dataset')\n",
    "    print(train_dataset)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        train_dataset, batch_size=256, shuffle=True,\n",
    "        num_workers=12, pin_memory=True, drop_last=True)\n",
    "\n",
    "    return train_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the quey and key encoder models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ISD(nn.Module):\n",
    "    def __init__(self, arch, K=65536, m=0.999, T=0.07):\n",
    "        super(ISD, self).__init__()\n",
    "\n",
    "        self.K = K\n",
    "        self.m = m\n",
    "        self.T = T\n",
    "\n",
    "        # create encoders and prediction layers\n",
    "        if 'resnet' in arch:\n",
    "            # both encoders should have same arch\n",
    "            self.encoder_q = resnet.__dict__[arch]()\n",
    "            self.encoder_k = resnet.__dict__[arch]()\n",
    "            # save output embedding dimensions\n",
    "            # assuming that both encoders have same dim\n",
    "            feat_dim = self.encoder_q.fc.in_features\n",
    "            out_dim = feat_dim\n",
    "\n",
    "            ##### prediction layer ####\n",
    "            # 1. have a prediction layer for q with BN\n",
    "            self.predict_q = nn.Sequential(\n",
    "                nn.Linear(feat_dim, feat_dim, bias=False),\n",
    "                nn.BatchNorm1d(feat_dim),\n",
    "                nn.ReLU(inplace=True),\n",
    "                nn.Linear(feat_dim, feat_dim, bias=True),\n",
    "            )\n",
    "\n",
    "            ##### projection layers ####\n",
    "            # 1. no projection layers for encoders\n",
    "            self.encoder_k.fc = nn.Sequential()\n",
    "            self.encoder_q.fc = nn.Sequential()\n",
    "        else:\n",
    "            raise ValueError('arch not found: {}'.format(arch))\n",
    "\n",
    "        # copy query encoder weights to key encoder\n",
    "        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n",
    "            param_k.data.copy_(param_q.data)\n",
    "            param_k.requires_grad = False\n",
    "\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def data_parallel(self):\n",
    "        self.encoder_q = torch.nn.DataParallel(self.encoder_q)\n",
    "        self.encoder_k = torch.nn.DataParallel(self.encoder_k)\n",
    "        self.predict_q = torch.nn.DataParallel(self.predict_q)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the loader and the encoders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==> train dataset\n",
      "Dataset ImageFolderEx\n",
      "    Number of datapoints: 1281167\n",
      "    Root location: /datasets/imagenet/train\n",
      "    StandardTransform\n",
      "Transform: <__main__.TwoCropsTransform object at 0x7f07eeae4210>\n"
     ]
    }
   ],
   "source": [
    "parser = parse_option()\n",
    "args = parser.parse_args([\n",
    "  '--num_workers', '16',\n",
    "  '/datasets/imagenet'\n",
    "])\n",
    "\n",
    "train_loader = get_train_loader(args)\n",
    "\n",
    "# memory bank size is 128k\n",
    "# momentum is 0.999\n",
    "# and temperature is 0.02\n",
    "isd = ISD('resnet18', K=128000, m=0.999, T=0.02)\n",
    "isd.data_parallel()\n",
    "isd = isd.cuda()\n",
    "\n",
    "# switch to eval mode\n",
    "isd.eval()\n",
    "for p in isd.parameters():\n",
    "    p.requires_grad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fill up the queue with keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# fill up the memory bank with keys\n",
    "queue = []\n",
    "\n",
    "for idx, (_, (im_q, im_k), _) in enumerate(train_loader):\n",
    "    im_q = im_q.cuda(non_blocking=True)\n",
    "    im_k = im_k.cuda(non_blocking=True)\n",
    "   \n",
    "    # calculate keys\n",
    "    with torch.no_grad():\n",
    "        k = isd.encoder_k(im_k)\n",
    "        k = nn.functional.normalize(k, dim=1)\n",
    "    \n",
    "    # fill up 128k samples in the queue\n",
    "    if (len(queue)*256) == 128000:\n",
    "        break\n",
    "\n",
    "    queue.append(k)\n",
    "    \n",
    "# concatenate all keys\n",
    "queue = torch.cat(queue, dim=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate the target (teacher) probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# calculate similarities\n",
    "sim_k = torch.mm(k, queue.t())\n",
    "\n",
    "# scale with temp\n",
    "sim_k /= isd.T\n",
    "\n",
    "# calculate probabilities\n",
    "targets = F.softmax(sim_k, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calcualte the output (student) probabilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the queries\n",
    "with torch.no_grad():\n",
    "    q = isd.encoder_q(im_q)\n",
    "    q = nn.functional.normalize(q, dim=1)\n",
    "\n",
    "# calculate similarities\n",
    "sim_q = torch.mm(q, queue.t())\n",
    "\n",
    "# scale with temp\n",
    "sim_q /= isd.T\n",
    "\n",
    "# calculate output log probabilities\n",
    "outputs = F.log_softmax(sim_q, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = F.kl_div(outputs, targets, reduction='batchmean')"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}