N-Tuplet Loss.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "# Plots\n",
    "%matplotlib inline\n",
    "from jupyterthemes import jtplot; jtplot.style(theme='monokai', context='notebook', ticks=True, grid=False);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "----------------------------------------------------------------\n        Layer (type)               Output Shape         Param #\n================================================================\n            Conv2d-1           [-1, 16, 32, 32]             448\n         MaxPool2d-2           [-1, 16, 32, 32]               0\n              ReLU-3           [-1, 16, 32, 32]               0\n       BatchNorm2d-4           [-1, 16, 32, 32]              32\n       HasherBlock-5           [-1, 16, 32, 32]               0\n            Conv2d-6           [-1, 20, 16, 16]           2,900\n         MaxPool2d-7           [-1, 20, 16, 16]               0\n              ReLU-8           [-1, 20, 16, 16]               0\n       BatchNorm2d-9           [-1, 20, 16, 16]              40\n      HasherBlock-10           [-1, 20, 16, 16]               0\n           Conv2d-11             [-1, 28, 8, 8]           5,068\n        MaxPool2d-12             [-1, 28, 8, 8]               0\n             ReLU-13             [-1, 28, 8, 8]               0\n      BatchNorm2d-14             [-1, 28, 8, 8]              56\n      HasherBlock-15             [-1, 28, 8, 8]               0\n           Conv2d-16             [-1, 32, 4, 4]           8,096\n        MaxPool2d-17             [-1, 32, 4, 4]               0\n             ReLU-18             [-1, 32, 4, 4]               0\n      BatchNorm2d-19             [-1, 32, 4, 4]              64\n      HasherBlock-20             [-1, 32, 4, 4]               0\n           Conv2d-21             [-1, 40, 2, 2]          11,560\n        MaxPool2d-22             [-1, 40, 2, 2]               0\n             ReLU-23             [-1, 40, 2, 2]               0\n      BatchNorm2d-24             [-1, 40, 2, 2]              80\n      HasherBlock-25             [-1, 40, 2, 2]               0\nAdaptiveMaxPool2d-26             [-1, 40, 1, 1]               0\n          Flatten-27                   [-1, 40]               0\n           Linear-28                   [-1, 24]             984\n             ReLU-29                   [-1, 24]               0\n           Linear-30                   [-1, 16]             400\n             Tanh-31                   [-1, 16]               0\n================================================================\nTotal params: 29,728\nTrainable params: 29,728\nNon-trainable params: 0\n----------------------------------------------------------------\nInput size (MB): 0.05\nForward/backward pass size (MB): 0.92\nParams size (MB): 0.11\nEstimated Total Size (MB): 1.08\n----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from Models import CustomHasher, SmallHasher, ResnetHasher\n",
    "SmallHasher().visualize()\n",
    "model = SmallHasher()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Training Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "class TripletDataset(torch.utils.data.Dataset):    \n",
    "\n",
    "    # Set max_images to overfit on small test set\n",
    "    def __init__(self, directory, transforms, max_images=np.inf):\n",
    "        self.transforms = transforms\n",
    "        \n",
    "        self.file_list = glob.glob(f'{directory}/*.png')\n",
    "        self.file_list.sort()\n",
    "        if (len(self.file_list) > max_images):\n",
    "            self.file_list = self.file_list[:max_images]\n",
    "        self.num_items_available = len(self.file_list)\n",
    "\n",
    "        # from imageaug\n",
    "        self._tt = ToRGBTensor()\n",
    "        \n",
    "        print(f'Found {len(self.file_list)} images.')\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_items_available\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        anchor   = Image.open(self.file_list[index])\n",
    "        positive = self.transforms(anchor)\n",
    "\n",
    "        return (self._tt(anchor), self._tt(positive))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image Augmentation/Transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from torchvision.transforms import Compose\n",
    "from imageaug import *\n",
    "\n",
    "transforms = Compose([\n",
    "    ApplyOne(RandomRotate(range(-10, 11, 1)), RandomFillCrop(1)),\n",
    "    ApplyOne(Noise(), Greyscale()),\n",
    "]) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Found 640 images.\nFound 0 images.\n"
     ]
    }
   ],
   "source": [
    "max_train = 64*10\n",
    "max_valid = 64*5\n",
    "dataset = TripletDataset('TrainDataset', transforms, max_images=max_train)\n",
    "validationset = TripletDataset('ValidDataset', transforms, max_images=max_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from vistools import show_imgpair\n",
    "\n",
    "show_imgpair(dataset[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "output_type": "error",
     "ename": "ValueError",
     "evalue": "num_samples should be a positive integer value, but got num_samples=0",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-2c9eaa780d5b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mDataLoader\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mtrain_loader\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataLoader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_workers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdrop_last\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mvalid_loader\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mDataLoader\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalidationset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnum_workers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdrop_last\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\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      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;31m# Model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, dataset, batch_size, shuffle, sampler, batch_sampler, num_workers, collate_fn, pin_memory, drop_last, timeout, worker_init_fn, multiprocessing_context, generator, prefetch_factor, persistent_workers)\u001b[0m\n\u001b[1;32m    260\u001b[0m                     \u001b[0;31m# Cannot statically verify that dataset is Sized\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;31m# Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]\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[0msampler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mRandomSampler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgenerator\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mgenerator\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\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    264\u001b[0m                     \u001b[0msampler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mSequentialSampler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/torch/utils/data/sampler.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, data_source, replacement, num_samples, generator)\u001b[0m\n\u001b[1;32m    102\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_samples\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_samples\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[1;32m    103\u001b[0m             raise ValueError(\"num_samples should be a positive integer \"\n\u001b[0;32m--> 104\u001b[0;31m                              \"value, but got num_samples={}\".format(self.num_samples))\n\u001b[0m\u001b[1;32m    105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    106\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mproperty\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: num_samples should be a positive integer value, but got num_samples=0"
     ]
    }
   ],
   "source": [
    "n = 64\n",
    "\n",
    "# Dataloader\n",
    "from torch.utils.data import DataLoader\n",
    "train_loader = DataLoader(dataset, num_workers=2, batch_size=n, shuffle=True, drop_last=True)\n",
    "valid_loader = DataLoader(validationset, num_workers=2, batch_size=n, shuffle=True, drop_last=True)\n",
    "\n",
    "# Model\n",
    "model = SmallHasher()\n",
    "model_output_dim = 16\n",
    "\n",
    "# Loss\n",
    "margin = 2\n",
    "p = 2\n",
    "\n",
    "# Optimizer\n",
    "opt = torch.optim.Adam(model.parameters(), lr=.01)\n",
    "\n",
    "# Export\n",
    "export_model = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Triplet Loss:\n",
    "$$L(\\displaystyle A,\\displaystyle P,\\displaystyle N,\\displaystyle \\alpha) = max\\Big( d(f(\\displaystyle A), f(\\displaystyle N)) − d(f(\\displaystyle A), f(\\displaystyle P)) + \\displaystyle \\alpha, 0 \\Big)$$\n",
    "\n",
    "Where:\n",
    "* ${\\displaystyle A}$ is an anchor input, \n",
    "* ${\\displaystyle P}$ is a positive input of the same class as ${\\displaystyle A}$, \n",
    "* ${\\displaystyle N}$ is a negative input of a different class from ${\\displaystyle A}$, \n",
    "* ${\\displaystyle \\alpha }$  is a margin between positive and negative pairs, \n",
    "* ${\\displaystyle f}$ is an embedding onto a metric space (the model we're training), and \n",
    "* $d$ is a distance function on that space. (In this case L1 or L2 norm, Manhattan and Euclidean distance respectively).\n",
    "\n",
    "Taken from FaceNet paper\n",
    "https://arxiv.org/abs/1503.03832"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N-Tuplet Loss\n",
    "Triplet loss as described above accomplishes unsupervised learning of embeddings. This is all well and good, but what if we wanted to up the ante a bit? Triplet loss pushes the embeddings of two images farther apart. But what if we wanted to push even more embeddings apart, probably more intelligently, in a single step? Behold my monstrosity, n-tuple loss."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let $f$ be an embedding from the input space onto metric space. In this case, $f$ represents our model.\n",
    "\n",
    "Let $d$ be a distance function on that metric space.\n",
    "\n",
    "Let $\\alpha$ be the desired distance margin between embeddings which do not have the same class.\n",
    "\n",
    "Let $A$ be a tuple of $n$ anchor inputs from separate classes, and $P$ be a tuple of $n$ corresponding positive matches from each class, and let that class be distinct. That is to say, $A_i$ and $P_i$ are of the same class, but the classes of $A_i$ and $A_j$, and of $A_i$ and $P_j$, are different for all $i,j$ from 0 to n-1 where $i \\neq j$.\n",
    "\n",
    "Finally, \n",
    "\n",
    "$$L(\\displaystyle A,\\displaystyle P,\\displaystyle \\alpha) = \\sum_{i=0}^{n-1} \\sum_{j=i+1}^{n-1} \\Big\\lbrack d(f(A_j), f(P_j)) - d(f(A_i),f(A_j)) + \\alpha \\Big\\rbrack_{+}.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## But, why?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main difference is that training the model in this way allows for more information to be digested from each batch. Triplet loss is effective, but it is wasteful in the sense that there's a lot of computation that could be done to spread the embeddings that isn't being done.\n",
    "\n",
    "Suppose that $n$ images are run through the model. Half of the images become anchor images, half become negative images, and positive images are generated through data augmentation from the anchors. For normal triplet loss, this is a batch of $\\frac{n}{2}$ comparisons. Although at the end of the forward pass we have embeddings in memory to compute the loss for an anchor/positive pair and an anchor, and the results would be just as meaningful as the one that triplet loss compares it to, we do not.\n",
    "\n",
    "However, if we run the same $n$ images through the model and do the above, we can do a lot more cross comparison. We let anchor images also act as negatives, generate the positives through data augmentation, and make every comparison possible. This ends up being ${n\\choose 2} = \\frac{n^2 - n}{2}$ comparisons."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of comparisons that can be done grows quadratically by the above rate. This results in many more comparisons than normal triplet loss, hopefully leading to a higher quality of gradient updates and more accurate convergence.\n",
    "\n",
    "Suppose $n = 64$, which seems like a reasonable sample size. With triplet loss, 32 comparisons would be done. This scheme does 2,016 of them, with no additional forward passes required. As the sample size increases, which tends to happen when people are trying to train models faster and faster, the difference gets even more pronounced. A sample size of 128 yeilds 8,128 comparisons compared to triplet loss's 64."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It should be noted however that because the magnitude of the weight update is determined linearly by the size of the batch, and the effort required to cross compare a large number of combinations of samples may become too expensive to justify for the increase in the quality of the gradient update, particularly at the start of training the model.\n",
    "\n",
    "There are many factors that go into optimization at the end. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N-Tuple Loss Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "2016\n"
     ]
    }
   ],
   "source": [
    "# Generate all possible combinations of (i, j) for n given by the above scheme\n",
    "from itertools import combinations\n",
    "ij_pairs = list(combinations(range(0, n), r=2))\n",
    "num_combinations = len(ij_pairs)\n",
    "print(num_combinations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Starting Trace.\n"
     ]
    },
    {
     "output_type": "error",
     "ename": "NameError",
     "evalue": "name 'model_output_dim' is not defined",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-15-d052581ebe86>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     25\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Starting Trace.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m     ntuplet_loss = torch.jit.trace(\n\u001b[0;32m---> 27\u001b[0;31m         ntuplet_loss, (torch.zeros([n, model_output_dim]), torch.zeros([n, model_output_dim])))\n\u001b[0m\u001b[1;32m     28\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Finished Trace.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'model_output_dim' is not defined"
     ]
    }
   ],
   "source": [
    "from torch.nn.functional import pairwise_distance\n",
    "\n",
    "def ntuplet_loss(anchor_embeddings, positive_embeddings):\n",
    "    # Embeddings are stacked tensors of batchnum many embeddings. These stacks have shape (n, embedding_size).\n",
    "\n",
    "    # Calculate distances between anchors and their respective positive matches.\n",
    "    ap_distances = pairwise_distance(anchor_embeddings, positive_embeddings, keepdim=True, p=p)\n",
    "\n",
    "    losses = []\n",
    "    for ij_pair in ij_pairs:\n",
    "        i, j = ij_pair\n",
    "        \n",
    "        combination_ij_distance = torch.dist(anchor_embeddings[i], anchor_embeddings[j], p=p)\n",
    "\n",
    "        # Minimizing this means maximizing combination_ij_distances and minimizing ap_distances.\n",
    "        dist = ap_distances[j] - combination_ij_distance + margin\n",
    "        \n",
    "        single_loss = torch.max(dist, torch.zeros_like(dist, requires_grad=True))\n",
    "        losses.append(single_loss)\n",
    "        \n",
    "    return sum(losses) / num_combinations\n",
    "\n",
    "trace = True\n",
    "if trace:\n",
    "    print('Starting Trace.')\n",
    "    ntuplet_loss = torch.jit.trace(\n",
    "        ntuplet_loss, (torch.zeros([n, model_output_dim]), torch.zeros([n, model_output_dim])))\n",
    "    print('Finished Trace.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Triplet Loss\n",
    "triplet_loss = torch.nn.TripletMarginLoss(margin=margin, p=p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "print_batches = True\n",
    "\n",
    "def forward_and_loss(batch, epoch):\n",
    "        # Train the first two batches with triplet loss, and ntuplet loss for the rest.\n",
    "        # if epoch <= -1:\n",
    "        ntuplet = False\n",
    "        if not ntuplet:\n",
    "                # Segment batch\n",
    "                anchor_images,   negative_images = torch.split(batch[0], batch[0].shape[0]//2, dim=0)\n",
    "                positive_images, unused          = torch.split(batch[1], batch[1].shape[0]//2, dim=0)\n",
    "                del unused\n",
    "\n",
    "                # Forward\n",
    "                anchor_embeddings   = model.forward(anchor_images)\n",
    "                positive_embeddings = model.forward(positive_images)\n",
    "                negative_embeddings = model.forward(negative_images)\n",
    "\n",
    "                # Backward\n",
    "                loss = triplet_loss(anchor_embeddings, positive_embeddings, negative_embeddings)\n",
    "                \n",
    "        else:\n",
    "                # Forward\n",
    "                anchor_embeddings   = model.forward(batch[0])\n",
    "                positive_embeddings = model.forward(batch[1])\n",
    "        \n",
    "                # Backward\n",
    "                loss = ntuplet_loss(anchor_embeddings, positive_embeddings)\n",
    "\n",
    "        return loss\n",
    "\n",
    "# Returns loss\n",
    "def train_batch(model, batch, batchnum):\n",
    "        batch_start_time = time()\n",
    "\n",
    "        # Enable Training\n",
    "        model.train()\n",
    "        model.zero_grad()\n",
    "\n",
    "        # Forward and calculate loss\n",
    "        loss = forward_and_loss(batch, epoch)\n",
    "\n",
    "        # Apply grad update\n",
    "        loss.backward()\n",
    "        opt.step()\n",
    "\n",
    "        # Print\n",
    "        batch_end_time = time()\n",
    "        loss_num = loss.item()\n",
    "        if print_batches:\n",
    "            print_batch('train', loss_num, dataset, batch_start_time, batch_end_time)\n",
    "        \n",
    "        return loss_num\n",
    "\n",
    "\n",
    "# Returns loss\n",
    "def valid_batch(model, vbatch, batch_num):\n",
    "        batch_start_time = time()\n",
    "\n",
    "        # Disable Training\n",
    "        with torch.no_grad():\n",
    "                model.eval()\n",
    "\n",
    "                loss = forward_and_loss(batch, epoch)\n",
    "\n",
    "                # Print\n",
    "                batch_end_time = time()\n",
    "                loss_num = loss.item()\n",
    "                if print_batches:\n",
    "                    print_batch('valid', loss_num, validationset, batch_start_time, batch_end_time)\n",
    "        \n",
    "        return loss_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_batch(tv, loss_num, dataset, batch_start_time, batch_end_time):\n",
    "        batch_time = round(batch_end_time-batch_start_time, 2)\n",
    "        num_batches = (dataset.num_items_available // n) - 1\n",
    "        est_epoch_time_seconds = batch_time*num_batches\n",
    "        est_epoch_time_minutes = round(est_epoch_time_seconds / 60, 2)\n",
    "        print(f'\\rCompleted {tv} batch: {str(batch_num).rjust(4, \" \")} of {num_batches} | \\\n",
    "Loss: {\"{:6.4f}\".format(loss_num)} | \\\n",
    "Time: {str(batch_time).rjust(6, \" \")} (est. {\"{:4.1f}\".format(est_epoch_time_minutes)} min for epoch)', end='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "best_weights = None\n",
    "best_loss = np.inf\n",
    "\n",
    "# Train for at most 1000 epochs. There's no way it will take that long.\n",
    "for epoch in range(1, 1000):\n",
    "    \n",
    "    # Training Batches\n",
    "    batch_num = 0\n",
    "    epoch_start_time = time()\n",
    "    for batch_num, batch in enumerate(train_loader):\n",
    "        train_batch(model, batch, batch_num)\n",
    "    \n",
    "    # So carriage return prints train/valid on different lines\n",
    "    if print_batches:\n",
    "        print()\n",
    "\n",
    "    # Validation batches\n",
    "    batch_num = 0\n",
    "    batch_valid_losses = []\n",
    "    for batch_num, vbatch in enumerate(valid_loader):\n",
    "        loss_num = valid_batch(model, vbatch, batch_num)\n",
    "        batch_valid_losses.append(loss_num)\n",
    "    valid_loss = sum(batch_valid_losses) / len(batch_valid_losses)\n",
    "    \n",
    "    # Print epoch results\n",
    "    epoch_end_time = time()\n",
    "    epoch_time_minutes = round((epoch_end_time - epoch_start_time) / 60, 2)\n",
    "    print(f\"\"\"\\n\\\n",
    "╔══════════════════════════════════════════════════════════════════╗\\n║ \\\n",
    "End of Epoch: {str(epoch).rjust(3, \" \")} | \\\n",
    "Validation Loss: {\"{:7.4f}\".format(valid_loss)} | \\\n",
    "Time: {str(epoch_time_minutes).rjust(6, \" \")} min. ║\\n\\\n",
    "╚══════════════════════════════════════════════════════════════════╝\\n\"\"\")\n",
    "\n",
    "    # Early stopping\n",
    "    # Cannot break the training loop on the first pass b/c best_loss is infinity.\n",
    "    if best_loss < valid_loss and epoch >= 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Export Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_path = 'Models/model%s.zip'\n",
    "\n",
    "\n",
    "def next_path(path_pattern):\n",
    "    i = 1\n",
    "    while os.path.exists(path_pattern % i):\n",
    "        i += 1\n",
    "    return path_pattern % i\n",
    "\n",
    "if export_model:\n",
    "    fname = next_path(model_path)\n",
    "    torch.save(model, fname)\n",
    "    print(f'Saved model as: {fname}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"Models/model2.zip\"\n",
    "model = torch.load(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt = ToRGBTensor()\n",
    "def load_img(fname):\n",
    "    return tt(Image.open(fname)).unsqueeze(0)\n",
    "def compare_vectors(v1, v2):\n",
    "    vec1, vec2 = v1[0], v2[0] # b/c second term is the file name\n",
    "    return np.linalg.norm(vec1-vec2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "embeddings = []\n",
    "\n",
    "# Turn off gradients\n",
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "\n",
    "    # Load all the validation data and hash it\n",
    "    for fname in tqdm(validationset.file_list):\n",
    "        embedding = model.forward(load_img(fname)).numpy()\n",
    "        embeddings.append((embedding, fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_vectors(embeddings[0], embeddings[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def show_matches(matches, axes=True):\n",
    "    fname_list = [tup[1][1] for tup in matches]\n",
    "    imgs = [Image.open(fname) for fname in fname_list]\n",
    "\n",
    "    axs = plt.subplots(1, len(imgs), figsize=(10, 5))[1]\n",
    "    axs = np.ravel(axs)\n",
    "\n",
    "    for i in range(len(imgs)):\n",
    "        axs[i].imshow(imgs[i])\n",
    "        if not axes:\n",
    "            axs[i].set_axis_off()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(embeddings))\n",
    "print(np.log2(len(embeddings)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Construct a vantage point tree for efficient nearest neighbor lookup\n",
    "from vptree import VPTree\n",
    "vpt = VPTree(embeddings, compare_vectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.imshow(Image.open('VP_Test_Img.png'));\n",
    "plt.show()\n",
    "\n",
    "show_matches(vpt.get_n_nearest_neighbors(embeddings[4], 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a vantage point tree for efficient nearest neighbor lookup\n",
    "from sys import setrecursionlimit\n",
    "setrecursionlimit(1500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torchvision.transforms.functional import rotate\n",
    "plt.imshow(rotate(Image.open(\"/home/apaz/FileNotFound.png\"), -30))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.3-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}