classification_and_metric_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from fastai.vision import *\n",
    "from fastai.metrics import accuracy\n",
    "from fastai.basic_data import *\n",
    "from skimage.util import montage\n",
    "import pandas as pd\n",
    "from torch import optim\n",
    "import re\n",
    "import pretrainedmodels\n",
    "from torch.nn.functional import cross_entropy\n",
    "\n",
    "import albumentations\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fixes ulimit issue: https://github.com/pytorch/pytorch/issues/973\n",
    "\n",
    "import resource\n",
    "rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)\n",
    "resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.46.dev0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import fastai\n",
    "fastai.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fastai\n",
    "from fastprogress import force_console_behavior\n",
    "import fastprogress\n",
    "fastprogress.fastprogress.NO_BAR = True\n",
    "master_bar, progress_bar = force_console_behavior()\n",
    "fastai.basic_train.master_bar, fastai.basic_train.progress_bar = master_bar, progress_bar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zen_dataset import *\n",
    "from zen_dataset.utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below follows a lot of code to set things up. I give an overview of how it all works together before I start to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "aug = albumentations.Compose([\n",
    "    albumentations.RandomBrightnessContrast(p=0.75),\n",
    "    albumentations.ShiftScaleRotate(shift_limit=0, scale_limit=0.1, rotate_limit=10, interpolation=2, p=0.75)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "composed_augs = albumentations.Compose([\n",
    "    albumentations.RandomBrightnessContrast(p=0.75),\n",
    "    albumentations.ShiftScaleRotate(shift_limit=0, scale_limit=0.1, rotate_limit=10, interpolation=2, p=0.75)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment(ary):\n",
    "    return composed_augs(image=ary)['image']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Reader():\n",
    "    def __init__(self, path, augment_fn=None):\n",
    "        self.path = path\n",
    "        self.augment_fn = augment_fn\n",
    "    def __call__(self, fns):\n",
    "        paths = [f'{self.path}/{filename}' for filename in fns]\n",
    "        images = [open_image(image_path) for image_path in paths]\n",
    "        tensors = [image2tensor(image, augment_fn = self.augment_fn) for image in images]\n",
    "        return [imagenet_normalize(tensor) for tensor in tensors]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Labeler():\n",
    "    def __init__(self):\n",
    "        df = pd.read_csv('data/train.csv')\n",
    "        self.fn2label = {}\n",
    "        for row in df[df.Id != 'new_whale'].itertuples():\n",
    "            self.fn2label[row.Image] = row.Id\n",
    "        self.classes = sorted(list(set(list(self.fn2label.values()))))\n",
    "    def __call__(self, fns):\n",
    "        labels = [self.fn2label[fn] for fn in fns]\n",
    "        return [self.classes.index(label) for label in labels] + [1 if labels[0] != labels[1] else 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.cnn = nn.Sequential(*list(models.resnet50(True).children())[:-2])\n",
    "        self.head = create_head(4096, 5004, [2048])\n",
    "        self.ada_concat = AdaptiveConcatPool2d(1)\n",
    "\n",
    "    def forward(self, ims_a, ims_b):\n",
    "        cnn_out_a = self.cnn(ims_a)\n",
    "        out_a = self.head(cnn_out_a)\n",
    "        \n",
    "        cnn_out_b = self.cnn(ims_b)\n",
    "        out_b = self.head(cnn_out_b)\n",
    "\n",
    "        return out_a, out_b, self.ada_concat(cnn_out_a).squeeze(), self.ada_concat(cnn_out_b).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://hackernoon.com/facial-similarity-with-siamese-networks-in-pytorch-9642aa9db2f7\n",
    "\n",
    "class ContrastiveLoss(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Contrastive loss function.\n",
    "    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, margin=2.0):\n",
    "        super(ContrastiveLoss, self).__init__()\n",
    "        self.margin = margin\n",
    "\n",
    "    def forward(self, output1, output2, label):\n",
    "        euclidean_distance = F.pairwise_distance(output1, output2)\n",
    "        loss_contrastive = torch.mean((1-label) * torch.pow(euclidean_distance, 2) +\n",
    "                                      (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))\n",
    "        return loss_contrastive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "MARGIN = 60\n",
    "\n",
    "def cross_entropy_loss(preds, labels_a, labels_b, diff_class_ind):\n",
    "    return F.cross_entropy(preds[0], labels_a) + F.cross_entropy(preds[1], labels_b)\n",
    "\n",
    "def contr_loss(preds, labels_a, labels_b, diff_class_ind):\n",
    "    c_loss = ContrastiveLoss(MARGIN)\n",
    "    return c_loss(preds[2], preds[3], diff_class_ind.float())\n",
    "\n",
    "def loss_fn(preds, labels_a, labels_b, diff_class_ind):\n",
    "    return 10 * cross_entropy_loss(preds, labels_a, labels_b, diff_class_ind) + contr_loss(preds, labels_a, labels_b, diff_class_ind) / 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy_mod(preds, labels_a, labels_b, diff_class_ind):\n",
    "    return 0.5 * accuracy(preds[0], labels_a) + 0.5 * accuracy(preds[1], labels_b)\n",
    "\n",
    "def map5_mod(preds, labels_a, labels_b, diff_class_ind):\n",
    "    return 0.5 * map5(preds[0], labels_a) + 0.5 * map5(preds[1], labels_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# When I refer to 'whale', I mean a particular image (the file name).\n",
    "\n",
    "df = pd.read_csv('data/train.csv')\n",
    "df = df[df.Id != 'new_whale']\n",
    "images_without_meaningful_bbox_predictions = \\\n",
    "    ['85a95e7a8.jpg', 'b370e1339.jpg', 'b4cb30afd.jpg', 'd4cb9d6e4.jpg', '6a72d84ca.jpg']\n",
    "df = df[~df.Image.isin(images_without_meaningful_bbox_predictions)]\n",
    "\n",
    "labeler = Labeler()\n",
    "\n",
    "def create_basic_dataloader(sz, batch_size, num_workers=12):\n",
    "    reader = Reader(f'data/train-extracted-{sz}')\n",
    "    basic_ds = Dataset([*zip(df.Image.tolist(), df.Image.tolist())], reader, labeler)\n",
    "    return DataLoader(basic_ds, batch_size=batch_size, num_workers=num_workers)\n",
    "\n",
    "def create_similarity_dict(model, dataloader):\n",
    "    # Calculating descriptors for each image\n",
    "    descs = []\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in dataloader:\n",
    "            ims = batch[0][0].cuda()\n",
    "            cnn_out = learn.model.cnn(ims)\n",
    "            descs.append(learn.model.ada_concat(cnn_out).squeeze().detach().cpu())\n",
    "\n",
    "    descs = torch.cat(descs).cuda()\n",
    "\n",
    "    # Calculating similarity dict for each image\n",
    "    dists = {}\n",
    "    for i, (whale, _) in enumerate(dataloader.items):\n",
    "        dists[whale] = torch.pairwise_distance(descs[i], descs).cpu().numpy()\n",
    "    \n",
    "    return dists\n",
    "\n",
    "def create_data(sz, dist_dict, batch_size, k=20, num_workers=12, train_on_both_train_and_val=False):\n",
    "    reader_aug = Reader(f'data/train-extracted-{sz}', augment_fn=augment)\n",
    "    reader = Reader(f'data/train-extracted-{sz}')\n",
    "    \n",
    "    val_fns = list(pd.read_pickle('data/val_fns'))\n",
    "    val_fns_set = set(val_fns)\n",
    "\n",
    "    trn_df = df[~df.Image.isin(val_fns)]\n",
    "    val_df = df[df.Image.isin(val_fns)]\n",
    "    \n",
    "    ds_on_which_dists_were_calculated = Dataset([*zip(df.Image.tolist(), df.Image.tolist())], reader, labeler)\n",
    "    \n",
    "    uniq_whales = df.Id.unique().tolist() if train_on_both_train_and_val else trn_df.Id.unique().tolist()\n",
    "\n",
    "    def sample_other_whale():\n",
    "        candidate_whales = dist_dict[this_whale].argsort() \n",
    "        this_whale_class = labeler.fn2label[this_whale]\n",
    "        candidate_fns = []\n",
    "        for i in range(200):\n",
    "            candidate_whale = ds_on_which_dists_were_calculated.items[candidate_whales[i]][0]\n",
    "            if (candidate_whale not in val_fns_set) and (labeler.fn2label[candidate_whale] != this_whale_class): \n",
    "                candidate_fns.append(candidate_whale)\n",
    "            if len(candidate_fns) == k: break \n",
    "        np.random.shuffle(candidate_fns) # randomly pick one from K toughest matches\n",
    "        return candidate_fns[0]\n",
    "\n",
    "    def sample_this_whale():\n",
    "        return this_whale_df.sample(n=1).iloc[0].Image\n",
    "\n",
    "    train_items = []\n",
    "    for whale in uniq_whales:\n",
    "        this_whale_df = trn_df[trn_df.Id == whale]\n",
    "        other_whale_df = trn_df[trn_df.Id != whale]\n",
    "\n",
    "        this_whale = sample_this_whale()\n",
    "\n",
    "        # sampling same whale if possible\n",
    "        if this_whale_df.shape[0] == 1: # only a single picture of this whale in dataset\n",
    "            other_whale = sample_other_whale()\n",
    "            train_items.append([this_whale, other_whale])\n",
    "        else:\n",
    "            same_whale = this_whale_df[this_whale_df.Image != this_whale].sample(n=1).iloc[0].Image\n",
    "            train_items.append([this_whale, same_whale])\n",
    "\n",
    "        # sampling different whales\n",
    "        this_whale = sample_this_whale()\n",
    "        train_items.append([this_whale, sample_other_whale()])\n",
    "    \n",
    "    if train_on_both_train_and_val:\n",
    "        valid_items = list(zip(val_df.Image.values[:batch_size].tolist(), val_df.Image.values[BS:2*batch_size].tolist()))\n",
    "    else:\n",
    "        valid_items = list(zip(val_df.Image.values[:1465].tolist(), val_df.Image.values[1465:2930].tolist()))\n",
    "\n",
    "    train_ds = Dataset(train_items, reader_aug, labeler)\n",
    "    valid_ds = Dataset(valid_items, reader, labeler)\n",
    "\n",
    "    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers)\n",
    "    valid_dl = DataLoader(valid_ds, batch_size=batch_size, num_workers=num_workers)\n",
    "\n",
    "    data = DataBunch(train_dl, valid_dl)\n",
    "    data.train_ds.loss_func = lambda: None\n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_fake_data(): # needed for loading the model\n",
    "    fake_ds = Dataset([],_,_)\n",
    "    fake_dl = DataLoader(fake_ds)\n",
    "\n",
    "    data = DataBunch(fake_dl, fake_dl)\n",
    "    data.train_ds.loss_func = lambda: None \n",
    "    \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.99 s, sys: 592 ms, total: 2.58 s\n",
      "Wall time: 2.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "learn = Learner(create_fake_data(), CustomModel(), loss_func=loss_fn, metrics=[accuracy_mod, map5_mod, cross_entropy_loss, contr_loss])\n",
    "learn = learn.clip_grad()\n",
    "learn.split((learn.model.cnn[6], learn.model.head))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I chose to implement a model based on resnet50 that would both classify each of the presented images as well as calculate dissimilarity between image pairs.\n",
    "\n",
    "Each training example consists of two images, most of them consisting of images of different whales and where possible of images of the same whale. I sample the images in a way as to maintain some class balance and to not favor whales with significantly more images.\n",
    "\n",
    "The model is presented with images A and B. It first sends the images through the convolution part of resnet50 (pretrained on imagenet). This way we obtain the 2048 feature maps of some dimensionality (the actual dimensionality of feature maps will depend on the size of the input). Once we have those, we run a classifier head on them to predict labels (whale ids) for each of the images.\n",
    "\n",
    "For each image pair the model outputs label prediction for image A, label prediction for image B, a 4096-length feature vector for image A and a 4096-length feature vector for image B.\n",
    "\n",
    "These outputs are then used for calculating the loss. I use a custom loss that combines cross entropy with contrastive loss.\n",
    "\n",
    "Below I generate initial data for the model to train on. Whale pairs are samples based on euclidean distance between the CNN features (after application of adaptive concatenation which doubles their lenght from 2048 to 4096). Controlling the `k` parameter is a proxy for how hard we want the sampled dataset to be. I also add some measure of randomness at multiple points to hopefully keep the datasets diverse while still balanced and challenging.\n",
    "\n",
    "In the later portions of the training I resample the dataset after each epoch.\n",
    "\n",
    "Classification output of a model trained in this fashion achieves around 0.86 on [the private LB](https://www.kaggle.com/c/humpback-whale-identification/leaderboard). Using similarity calculations solely, the performance improves to 0.9.\n",
    "\n",
    "The training procedure looked as followed:\n",
    "* train the classification model on extracted bounding boxes without dataset construction as in earlier notebooks\n",
    "* load the weights into the custom model (modifying the state dict and loading of weights which is not shown here)\n",
    "* train on 224x224 images as below\n",
    "* train on 448x448 images\n",
    "* train on 448x448 images only with contrastive loss\n",
    "\n",
    "For the later parts of the training I generated new datasets every epoch with K as low as 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 23s, sys: 27.7 s, total: 1min 51s\n",
      "Wall time: 1min 51s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "SZ = 224\n",
    "NUM_WORKERS = 12\n",
    "BS = 32\n",
    "\n",
    "basic_dataloader = create_basic_dataloader(SZ, BS, NUM_WORKERS)\n",
    "dists = create_similarity_dict(learn.model, basic_dataloader)\n",
    "data = create_data(SZ, dists, BS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I am training from scratch. I first train the classifier head with the rest of the model frozen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss  accuracy_mod  map5_mod  cross_entropy_loss  contr_loss\n",
      "1         148.910690  147.522888  0.055290      0.076900  14.646444           26.461220   \n",
      "2         123.472168  145.344574  0.082594      0.108106  14.397060           34.348557   \n",
      "3         85.668175   135.626740  0.125256      0.165592  13.510564           13.028355   \n",
      "4         53.227573   127.417381  0.187713      0.246024  12.650944           22.698233   \n",
      "5         37.310898   117.722824  0.256655      0.318282  11.670763           25.379938   \n",
      "6         30.048853   112.574509  0.309215      0.374215  11.157745           24.926008   \n",
      "7         22.527187   108.538574  0.356655      0.419767  10.791218           15.660150   \n",
      "8         17.551012   99.792419   0.403754      0.465836  9.885334            23.477522   \n",
      "9         14.864764   91.965767   0.449829      0.508680  9.101912            23.665983   \n",
      "10        12.320655   83.247070   0.477816      0.538737  8.249646            18.765621   \n",
      "11        11.685419   78.960571   0.494198      0.554721  7.826274            17.445515   \n",
      "12        11.117297   78.526299   0.499659      0.558817  7.778657            18.493265   \n",
      "Total time: 19:22\n",
      "xbcfpkbgev\n",
      "CPU times: user 13min 24s, sys: 5min 40s, total: 19min 5s\n",
      "Wall time: 19min 23s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "learn = Learner(data, CustomModel(), loss_func=loss_fn, metrics=[accuracy_mod, map5_mod, cross_entropy_loss, contr_loss])\n",
    "learn = learn.clip_grad()\n",
    "learn.split((learn.model.cnn[6], learn.model.head))\n",
    "learn.freeze()\n",
    "\n",
    "learn.fit_one_cycle(12, 1e-2)\n",
    "learn.save(name())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I then proceed to training the entire model. I use the one cycle policy and use discriminative fine tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss  accuracy_mod  map5_mod  cross_entropy_loss  contr_loss\n",
      "1         26.731331   74.631279   0.519795      0.576769  7.457275            1.462620    \n",
      "2         21.375389   75.528976   0.516724      0.571519  7.549324            0.893103    \n",
      "3         20.266861   76.560631   0.497611      0.556661  7.649171            1.723411    \n",
      "4         18.548521   77.408104   0.483959      0.544317  7.732693            2.029275    \n",
      "5         17.174431   78.717873   0.479181      0.537947  7.857864            3.480946    \n",
      "6         16.853291   72.678078   0.521160      0.577253  7.253875            3.482910    \n",
      "7         15.177959   75.314728   0.501024      0.559454  7.515693            3.944957    \n",
      "8         13.921166   70.167458   0.518430      0.578083  7.001643            3.775323    \n",
      "9         13.039845   67.782875   0.534471      0.590421  6.761198            4.272540    \n",
      "10        10.834149   66.399384   0.544369      0.597173  6.614733            6.301155    \n",
      "11        10.516848   63.929035   0.555631      0.609278  6.361138            7.941088    \n",
      "12        8.837101    61.644421   0.564164      0.618168  6.137127            6.828721    \n",
      "13        8.199161    61.322220   0.576792      0.627218  6.088045            11.044115   \n",
      "14        7.196623    60.468719   0.575768      0.630205  5.998930            11.985427   \n",
      "15        6.209642    60.272873   0.570648      0.627389  5.976564            12.680258   \n",
      "16        5.264899    59.344940   0.577133      0.632827  5.891311            10.795804   \n",
      "17        4.952946    59.042267   0.583276      0.636030  5.851012            13.303552   \n",
      "18        4.683704    58.793663   0.582253      0.635950  5.826302            13.265815   \n",
      "19        4.361069    58.649273   0.581911      0.635421  5.803203            15.430987   \n",
      "20        4.349578    59.054955   0.583276      0.635921  5.844341            15.288428   \n",
      "Total time: 42:51\n",
      "gabgottwbr\n"
     ]
    }
   ],
   "source": [
    "learn.unfreeze()\n",
    "dists = create_similarity_dict(learn.model, basic_dataloader)\n",
    "learn.data = create_data(SZ, dists, BS)\n",
    "\n",
    "max_lr = 1e-3\n",
    "lrs = [max_lr/100, max_lr/10, max_lr]\n",
    "\n",
    "learn.fit_one_cycle(20, lrs)\n",
    "learn.save(name())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I used the results on the validation set to understand what effect changes had on performance, to pick hyperparameters, etc. For this competition, the validation set removed a lot of valuable information from the train set.\n",
    "\n",
    "As such, to complete the training, I switch to training on the entire train set (without retaining any images for the validation set).\n",
    "\n",
    "(I could do that through the insights I gained earlier and also because I knew my model would not overfit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.metrics = [] # metrics calculated on the validation set will no longer be useful"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss\n",
      "1         12.779222   62.867638   \n",
      "Total time: 02:02\n",
      "epoch     train_loss  valid_loss\n",
      "1         12.693471   59.829559   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         12.260114   61.417843   \n",
      "Total time: 02:02\n",
      "epoch     train_loss  valid_loss\n",
      "1         12.241240   63.250805   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         12.103703   65.692261   \n",
      "Total time: 02:03\n",
      "Finished training with lr: 0.0005\n",
      "fmnmnrqjyx\n"
     ]
    }
   ],
   "source": [
    "max_lr = 5e-4\n",
    "lrs = [max_lr/100, max_lr/10, max_lr]\n",
    "\n",
    "for _ in range(5):\n",
    "    dists = create_similarity_dict(learn.model, basic_dataloader)\n",
    "    learn.data = create_data(SZ, dists, BS, k=10, train_on_both_train_and_val=True)\n",
    "    learn.fit(1, lrs)\n",
    "print(f'Finished training with lr: {max_lr}')\n",
    "\n",
    "learn.save(name())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The train loss is much higher, but the sampling of whales here is harder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss\n",
      "1         11.932597   61.330284   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         11.082793   59.027733   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         10.894583   57.672829   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         11.322955   58.176491   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         10.782765   57.838867   \n",
      "Total time: 02:03\n",
      "Finished training with lr: 0.0001\n",
      "azhovmyckz\n"
     ]
    }
   ],
   "source": [
    "max_lr = 1e-4\n",
    "lrs = [max_lr/100, max_lr/10, max_lr]\n",
    "\n",
    "for _ in range(5):\n",
    "    dists = create_similarity_dict(learn.model, basic_dataloader)\n",
    "    learn.data = create_data(SZ, dists, BS, k=7, train_on_both_train_and_val=True)\n",
    "    learn.fit(1, lrs)\n",
    "print(f'Finished training with lr: {max_lr}')\n",
    "\n",
    "learn.save(name())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch     train_loss  valid_loss\n",
      "1         10.938106   57.279255   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         10.535534   56.972042   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         10.983456   56.210239   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         10.421021   57.603001   \n",
      "Total time: 02:03\n",
      "epoch     train_loss  valid_loss\n",
      "1         11.137517   57.158092   \n",
      "Total time: 02:03\n",
      "Finished training with lr: 5e-05\n",
      "nfbjfylcqh\n"
     ]
    }
   ],
   "source": [
    "max_lr = 5e-5\n",
    "lrs = [max_lr/100, max_lr/10, max_lr]\n",
    "\n",
    "for _ in range(5):\n",
    "    dists = create_similarity_dict(learn.model, basic_dataloader)\n",
    "    learn.data = create_data(SZ, dists, BS, k=3, train_on_both_train_and_val=True)\n",
    "    learn.fit(1, lrs)\n",
    "print(f'Finished training with lr: {max_lr}')\n",
    "\n",
    "learn.save(name())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is not exactly how I trained, but it hopefully captures the gist of it in a readable way.\n",
    "\n",
    "At this point I would switch to training on 448x448 crops. I would train on 448x448 crops with the custom loss function (combining cross entropy and contrastive divergence). For the last segment of the training, I used contrastive loss only.\n",
    "\n",
    "I will not carry out the training on larger images here, let us rather proceed to generating a submission."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Reader():\n",
    "    def __call__(self, paths):\n",
    "        images = [open_image(image_path) for image_path in paths]\n",
    "        tensors = [image2tensor(image) for image in images]\n",
    "        return [imagenet_normalize(tensor) for tensor in tensors]\n",
    "\n",
    "train_items = df.Image.apply(lambda fn: f'data/train-extracted-{SZ}/{fn}').tolist()\n",
    "test_items = list(map(lambda p: str(p), paths_to_files_in(f'data/test-extracted-{SZ}')))\n",
    "\n",
    "train_and_test_items = train_items + test_items\n",
    "\n",
    "train_and_test_ds = Dataset([*zip(train_and_test_items)], Reader(), lambda _: 0)\n",
    "train_and_test_dl = DataLoader(train_and_test_ds, batch_size=BS, num_workers=NUM_WORKERS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I run the cnn part of the model on all the images in the train and test sets to obtain features (I call them descs for descriptors)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 27.8 s, sys: 14 s, total: 41.8 s\n",
      "Wall time: 42.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "descs = []\n",
    "learn.model.eval()\n",
    "with torch.no_grad():\n",
    "    for batch in train_and_test_dl:\n",
    "        ims = batch[0][0].cuda()\n",
    "        cnn_out = learn.model.cnn(ims)\n",
    "        descs.append(learn.model.ada_concat(cnn_out).squeeze().detach().cpu())\n",
    "\n",
    "descs = torch.cat(descs).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I calculate distances between each image in the test set and all images in the train set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "def path2whale(path):\n",
    "    return re.search('(\\w*.\\w*$)', path).group(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 17.7 s, sys: 9.37 s, total: 27.1 s\n",
      "Wall time: 27.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "dists = {}\n",
    "for i, path in enumerate(train_and_test_dl.items[15694:]):\n",
    "    whale = path2whale(path[0])\n",
    "    dists[whale] = torch.pairwise_distance(descs[i + 15694], descs[:15694]).cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And I generate whale id predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_fns = [p.name for p in paths_to_files_in(f'data/test-extracted-{SZ}')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 17.1 s, sys: 4 ms, total: 17.1 s\n",
      "Wall time: 17.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "new_whale_threshold = 47\n",
    "\n",
    "all_preds = []\n",
    "for fn in test_fns:\n",
    "    most_similar = list(dists[fn].argsort())\n",
    "    preds = []\n",
    "    \n",
    "    while len(preds) < 5:\n",
    "        similar = most_similar.pop(0)\n",
    "        class_of_similar = labeler.fn2label[path2whale(train_and_test_items[similar])]\n",
    "        if dists[fn][similar] > new_whale_threshold:\n",
    "            if 'new_whale' not in preds: preds.append('new_whale')\n",
    "        if len(preds) < 5:\n",
    "            if class_of_similar not in preds: preds.append(class_of_similar)\n",
    "    all_preds.append(preds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I tried looking for the best threshold for predicting `new_whale` in a couple of ways. Based on thinking about this and the results I was seeing I came to the conclusion that predicting `new_whale` as first prediction a little more often than 27% of the time (which was the ratio of new whales in the public portion of the test set) should work quite well.\n",
    "\n",
    "One can alter this by modifying the `new_whale_threshold` above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3466515893956527"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean([preds[0] == 'new_whale' for preds in all_preds])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's generate the submission."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_name = 'res50_similarity'\n",
    "\n",
    "sub = pd.DataFrame({'Image': test_fns, 'Id': all_preds})\n",
    "sub.Id = sub.Id.str.join(' ')\n",
    "sub.to_csv(f'subs/{sub_name}.csv.gz', index=False, compression='gzip')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There was an image missing from the test set (one where I was unable to extract a bounding box) so here I am adding a prediction for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub = pd.read_csv(f'subs/{sub_name}.csv.gz')\n",
    "sub.append({'Image': '6a72d84ca.jpg', 'Id': 'new_whale'}, ignore_index=True).to_csv(f'subs/{sub_name}.csv.gz', index=False, compression='gzip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████| 186k/186k [00:01<00:00, 73.7kB/s]\n",
      "Successfully submitted to Humpback Whale Identification"
     ]
    }
   ],
   "source": [
    "!kaggle competitions submit -c humpback-whale-identification -f subs/{sub_name}.csv.gz -m \"{sub_name}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model as trained above achieves 0.84812 on private LB. With a bit more training on 448x448 images the score increased to 0.90813. As I have not spent a lot of time training the model, there is some chance the score would improve further with more training."
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}