Module13/13_siamese_triplet_mnist_empty.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "13_siamese_triplet_mnist.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "53--KowM64OE",
        "azGVIKYV1E0e"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ei8ToiWt1EzS"
      },
      "source": [
        "# Siamese networks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kd38ctpK1EzT"
      },
      "source": [
        "## Colab preparation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sx9e_pXlCuti"
      },
      "source": [
        "%load_ext autoreload\n",
        "%autoreload 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aU__26HN1EzZ"
      },
      "source": [
        "from os import path\n",
        "\n",
        "import numpy as np\n",
        "import random\n",
        "\n",
        "%matplotlib inline\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "from torch.optim import lr_scheduler\n",
        "import torch.optim as optim\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torchvision.datasets import MNIST\n",
        "from torchvision import transforms\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UMMut8UVCutt"
      },
      "source": [
        "# 1. Setup and initializations\n",
        "We'll go through learning feature embeddings using different loss functions on MNIST dataset. This is just for visualization purposes, thus we'll be using 2-dimensional embeddings which isn't the best choice in practice.\n",
        "\n",
        "For every experiment the same embedding network is used (`32 conv 5x5 -> ReLU -> MaxPool 2x2 -> 64 conv 5x5 -> ReLU -> MaxPool 2x2 -> Fully Connected 256 -> ReLU -> Fully Connected 256 -> ReLU -> Fully Connected 2`) with the same hyperparameters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tvM4mzUm1Ezc"
      },
      "source": [
        "class ExperimentParams():\n",
        "    def __init__(self):\n",
        "        self.num_classes = 10\n",
        "        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "        self.batch_size = 256\n",
        "        self.lr = 1e-2\n",
        "        self.num_epochs = 10\n",
        "        self.num_workers = 4\n",
        "        self.data_dir = '/home'\n",
        "        \n",
        "\n",
        "args = ExperimentParams()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BcmGBqXeCutw"
      },
      "source": [
        "## 1.1 Prepare dataset\n",
        "We'll be working on MNIST dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AQSHPH9P0BNx"
      },
      "source": [
        "\n",
        "mean, std = 0.1307, 0.3081\n",
        "\n",
        "train_dataset = MNIST(f'{args.data_dir}/data/MNIST', train=True, download=True,\n",
        "                             transform=transforms.Compose([\n",
        "                                 transforms.ToTensor(),\n",
        "                                 transforms.Normalize((mean,), (std,))\n",
        "                             ]))\n",
        "\n",
        "test_dataset = MNIST(f'{args.data_dir}/data/MNIST', train=False, download=True,\n",
        "                            transform=transforms.Compose([\n",
        "                                transforms.ToTensor(),\n",
        "                                transforms.Normalize((mean,), (std,))\n",
        "                            ]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TcZTFRnjCut3"
      },
      "source": [
        "\r\n",
        "## 1.2 Common setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dz2xh66UCut5"
      },
      "source": [
        "\n",
        "mnist_classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']\n",
        "colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728',\n",
        "              '#9467bd', '#8c564b', '#e377c2', '#7f7f7f',\n",
        "              '#bcbd22', '#17becf']\n",
        "\n",
        "def plot_embeddings(embeddings, targets, title='',xlim=None, ylim=None):\n",
        "    plt.figure(figsize=(10,10))\n",
        "    for i in range(10):\n",
        "        inds = np.where(targets==i)[0]\n",
        "        plt.scatter(embeddings[inds,0], embeddings[inds,1], alpha=0.5, color=colors[i])\n",
        "    if xlim:\n",
        "        plt.xlim(xlim[0], xlim[1])\n",
        "    if ylim:\n",
        "        plt.ylim(ylim[0], ylim[1])\n",
        "    plt.legend(mnist_classes)\n",
        "    plt.title(title)\n",
        "\n",
        "def extract_embeddings(dataloader, model, args):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        embeddings = np.zeros((len(dataloader.dataset), 2))\n",
        "        labels = np.zeros(len(dataloader.dataset))\n",
        "        k = 0\n",
        "        for images, target in dataloader:\n",
        "            images = images.to(args.device)\n",
        "            embeddings[k:k+len(images)] = model.get_embedding(images).data.cpu().numpy()\n",
        "            labels[k:k+len(images)] = target.numpy()\n",
        "            k += len(images)\n",
        "    return embeddings, labels\n",
        "\n",
        "\n",
        "def get_raw_images(dataloader,mean=0.1307, std=0.3081):\n",
        "\n",
        "    raw_images = np.zeros((len(dataloader.dataset), 1, 28, 28))\n",
        "    k = 0\n",
        "    for input, target in dataloader:\n",
        "        raw_images[k:k+len(input)] = (input*std + mean).data.cpu().numpy()\n",
        "        k += len(input)\n",
        "\n",
        "    return raw_images\n",
        "\n",
        "\n",
        "def show(img, title=None):\n",
        "    # img is a torch.Tensor     \n",
        "    npimg = img.numpy()\n",
        "    plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')\n",
        "    plt.axis('off')\n",
        "    if title is not None:\n",
        "        plt.title(title)\n",
        "    plt.pause(0.001)  # pause a bit so that plots are updated"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "75moY8AyCut_"
      },
      "source": [
        "# 2. Baseline: Classification with softmax\n",
        "We'll train the model for classification and use outputs of penultimate layer as embeddings. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJCmiKwf1Ezl"
      },
      "source": [
        "We will define our base embedding architecture which will serve as common backbone for our experiments"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ew1qgU_G1Ezl"
      },
      "source": [
        "## 2.1 Architecture"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49bn26xC1Ezm"
      },
      "source": [
        "### Exercise\n",
        "\n",
        "Complete the missing blocks in the definition of the following `EmbeddingNet` architecture: (`32 conv 5x5 -> ReLU -> MaxPool 2x2 -> 64 conv 5x5 -> ReLU -> MaxPool 2x2 -> Fully Connected 256 -> ReLU -> Fully Connected 256 -> ReLU -> Fully Connected 2`)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C1LjKH_21Ezn"
      },
      "source": [
        "\n",
        "class EmbeddingNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(EmbeddingNet, self).__init__()\n",
        "#         self.conv1 = nn.Conv2d(1, ...)\n",
        "#         self.conv2 = ...\n",
        "#         self.fc1 = ...\n",
        "#         self.fc2 = ...\n",
        "#         self.fc3 = ...\n",
        "\n",
        "    def forward(self, x, debug=False):\n",
        "        x1 = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=2, stride=2)\n",
        "#       output = ...                  \n",
        "        if debug == True:\n",
        "            print(f'input: {x.size()}')\n",
        "            print(f'x1: {x1.size()}')\n",
        "                      \n",
        "        return output\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.forward(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lOYKYEm41Ezq"
      },
      "source": [
        "If you want to better check the sizes of the hidden states and do debugging, you can add a `debug` variable in the `forward` function just like above"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D1PZbEVB1Ezr"
      },
      "source": [
        "input = torch.zeros(1, 1, 28, 28)\n",
        "net = EmbeddingNet()\n",
        "output = net(input,debug=True)\n",
        "print(f'output size: {output.size()}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IzMHlrtK6pxC"
      },
      "source": [
        "### Question\r\n",
        "The dimension of the output is `batch-size x 2`. Why?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uGMMtw_F1Ezv"
      },
      "source": [
        "Now let's define a classification net that will add fully connected layer on top of `EmbeddingNet`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JpWMG8kF1Ezv"
      },
      "source": [
        "### Exercise\n",
        "\n",
        "Fill in the missing spots in the `forward` pass. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7ChpUMXh1Ezw"
      },
      "source": [
        "\n",
        "class ClassificationNet(nn.Module):\n",
        "    def __init__(self, embedding_net, num_classes):\n",
        "        super(ClassificationNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "        self.fc = nn.Linear(2, num_classes)\n",
        "\n",
        "    def forward(self, x, debug=False):\n",
        "        # replace None with necessary entry \n",
        "        embedding = None\n",
        "        output = F.softplus(self.fc(embedding))\n",
        "        \n",
        "        # if debug == True:\n",
        "        #     print(f'input: {x.size()}')\n",
        "        #     print(f'embedding: {embedding.size()}')\n",
        "        #     print(f'output: {output.size()}')\n",
        "            \n",
        "        return output\n",
        "    \n",
        "    def get_embedding(self, x):\n",
        "        # replace None with necessary entry \n",
        "        return None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e5ib15cj1Ezy"
      },
      "source": [
        "## 2.2 Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KA8szJpj1Ezy"
      },
      "source": [
        "# Set up data loaders\n",
        "\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True} \n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "embedding_net = EmbeddingNet()\n",
        "model = ClassificationNet(embedding_net, num_classes=args.num_classes)\n",
        "loss_fn = torch.nn.CrossEntropyLoss()\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i_NBcq3w1Ez3"
      },
      "source": [
        "train_embeddings_baseline, train_labels_baseline = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_baseline, train_labels_baseline, 'Train embeddings before training')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eOLUpe-U1Ez6"
      },
      "source": [
        "def train_classif_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=50):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, total_corrects, num_samples = 0, 0, 0\n",
        "    corrects = 0    \n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        num_samples += data.size(0)\n",
        "        \n",
        "        data, target = data.to(args.device), target.to(args.device)\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(data)\n",
        "\n",
        "        loss = loss_fn(outputs, target)\n",
        "        losses.append(loss.data.item())\n",
        "\n",
        "        _,preds = torch.max(outputs.data,1)\n",
        "        corrects += torch.sum(preds == target.data).cpu()\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train: [{}/{} ({:.0f}%)]\\tLoss: {:.6f} \\tAccuracy: {}'.format(\n",
        "                batch_idx * len(data[0]), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), np.mean(losses), float(total_corrects)/num_samples))           \n",
        "            \n",
        "            total_loss += np.sum(losses)\n",
        "            total_corrects += corrects\n",
        "            losses, corrects = [], 0\n",
        "\n",
        "    return total_loss/(batch_idx + 1), total_corrects/num_samples\n",
        "\n",
        "def test_classif_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        losses, corrects = [], 0\n",
        "        num_samples = 0\n",
        "    \n",
        "        for batch_idx, (data, target) in enumerate(test_loader):\n",
        "\n",
        "            num_samples += data.size(0)\n",
        "            data, target = data.to(args.device), target.to(args.device)\n",
        "\n",
        "            outputs = model(data)\n",
        "\n",
        "            loss = loss_fn(outputs, target)\n",
        "            losses.append(loss.data.item())\n",
        "\n",
        "            _,preds = torch.max(outputs.data,1)\n",
        "            corrects += torch.sum(preds == target.data).cpu()\n",
        "\n",
        "        return np.sum(losses)/(batch_idx + 1), corrects/num_samples"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VakGpqxP6ReQ"
      },
      "source": [
        "### Question\r\n",
        "Why do we need `optimizer.zero_grad()`? What happens if we remove it?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yfGNHVAz1Ez9"
      },
      "source": [
        "start_epoch = 0\n",
        "\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "\n",
        "    train_loss, train_accuracy = train_classif_epoch(train_loader, model, loss_fn, optimizer, args)\n",
        "\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f} Average accuracy: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss, train_accuracy)\n",
        "    \n",
        "    val_loss, val_accuracy = test_classif_epoch(test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}  Average accuracy: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             val_loss, val_accuracy)\n",
        "    print(message)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OIjXt_931E0A"
      },
      "source": [
        "## 2.3 Visualizations\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q9o-6-ytCuuT"
      },
      "source": [
        "train_embeddings_baseline, train_labels_baseline = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_baseline, train_labels_baseline, 'Train embeddings classification')\n",
        "test_embeddings_baseline, test_labels_baseline = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_baseline, test_labels_baseline, 'Test embeddings classification')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KCL_UfwRSUii"
      },
      "source": [
        "train_embeddings_baseline, train_labels_baseline = extract_embeddings(train_loader, model, args)\r\n",
        "plot_embeddings(train_embeddings_baseline, train_labels_baseline, 'Train embeddings classification')\r\n",
        "test_embeddings_baseline, test_labels_baseline = extract_embeddings(test_loader, model, args)\r\n",
        "plot_embeddings(test_embeddings_baseline, test_labels_baseline, 'Test embeddings classification')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fDdsCIiHCuub"
      },
      "source": [
        "While the embeddings look separable (which is what we trained them for), they don't have good metric properties. They might not be the best choice as a descriptor for new classes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j9FhHE-tCuuc"
      },
      "source": [
        "# 3. Siamese network\n",
        "Now we'll train a siamese network that takes a pair of images and trains the embeddings so that the distance between them is minimized if their from the same class or greater than some margin value if they represent different classes.\n",
        "We'll minimize a contrastive loss function*:\n",
        "$$L_{contrastive}(x_0, x_1, y) = \\frac{1}{2} y \\lVert f(x_0)-f(x_1)\\rVert_2^2 + \\frac{1}{2}(1-y)\\{max(0, m-\\lVert f(x_0)-f(x_1)\\rVert_2)\\}^2$$\n",
        "\n",
        "*Raia Hadsell, Sumit Chopra, Yann LeCun, [Dimensionality reduction by learning an invariant mapping](http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf), CVPR 2006*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GfA7R5dF1E0F"
      },
      "source": [
        "## 3.1 Architecture\n",
        "We will first define the siamese architecture on top of our `EmbeddingNet`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNEzScHM1E0F"
      },
      "source": [
        "### Exercise\n",
        "\n",
        "Fill in the forward part of `SiameseNet`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hj7Jzzwz1E0G"
      },
      "source": [
        "class SiameseNet(nn.Module):\n",
        "    def __init__(self, embedding_net):\n",
        "        super(SiameseNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "\n",
        "    def forward(self, x1, x2):\n",
        "        # fill in the missing 2 lines :)\n",
        "        output1 = self.embedding_net(x1)\n",
        "        output2 = self.embedding_net(x2)\n",
        "        \n",
        "        return output1, output2\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.embedding_net(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "935Hwn0R1E0J"
      },
      "source": [
        "## 3.2 Data loader\n",
        "We will also need to adapt our data loader to fetch pairs of images "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XAFAu0We1E0K"
      },
      "source": [
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data.sampler import BatchSampler\n",
        "from PIL import Image\n",
        "\n",
        "class SiameseMNIST(Dataset):\n",
        "    \"\"\"\n",
        "    train mode: For each sample creates randomly a positive or a negative pair\n",
        "    test mode: Creates fixed pairs for testing\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, mnist_dataset):\n",
        "        self.mnist_dataset = mnist_dataset\n",
        "\n",
        "        self.train = self.mnist_dataset.train\n",
        "        self.transform = self.mnist_dataset.transform\n",
        "\n",
        "        if self.train:\n",
        "            self.train_labels = self.mnist_dataset.train_labels\n",
        "            self.train_data = self.mnist_dataset.train_data\n",
        "            self.labels_set = set(self.train_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.train_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "        else:\n",
        "            # generate fixed pairs for testing\n",
        "            self.test_labels = self.mnist_dataset.test_labels\n",
        "            self.test_data = self.mnist_dataset.test_data\n",
        "            self.labels_set = set(self.test_labels.numpy())\n",
        "            '''\n",
        "            create a dictionary with an entry key for each label and the value an array storing\n",
        "            the indices of the images having the respective label\n",
        "            '''\n",
        "            self.label_to_indices = {label: np.where(self.test_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "            random_state = np.random.RandomState(42)\n",
        "            # itereate through test_data and randomly select samples with the same label\n",
        "            positive_pairs = [[i,\n",
        "                               random_state.choice(self.label_to_indices[self.test_labels[i].item()]),\n",
        "                               1]\n",
        "                              for i in range(0, len(self.test_data), 2)]\n",
        "\n",
        "            # itereate through test_data, create a list of all labels different from current one and then\n",
        "            # randomly select samples with having one of these labels\n",
        "            negative_pairs = [[i,\n",
        "                               random_state.choice(self.label_to_indices[\n",
        "                                                       np.random.choice(\n",
        "                                                           list(self.labels_set - set([self.test_labels[i].item()]))\n",
        "                                                       )\n",
        "                                                   ]),\n",
        "                               0]\n",
        "                              for i in range(1, len(self.test_data), 2)]\n",
        "            # format: [index1, index2, label(0/1)]\n",
        "            self.test_pairs = positive_pairs + negative_pairs\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        \n",
        "        # at train time pairs of samples are fetched randomly on the fly\n",
        "        if self.train:\n",
        "            # select random label,i.e. similar (1) or non-similar (0) images\n",
        "            target = np.random.randint(0, 2)\n",
        "            img1, label1 = self.train_data[index], self.train_labels[index].item()\n",
        "            if target == 1:\n",
        "                # select an image with the same label as img1\n",
        "                siamese_index = index\n",
        "                while siamese_index == index:\n",
        "                    siamese_index = np.random.choice(self.label_to_indices[label1])\n",
        "            else:\n",
        "                # eliminate label1 from the set of possible labels to select\n",
        "                siamese_label = np.random.choice(list(self.labels_set - set([label1])))\n",
        "                # randomly select an image having a label from this subset\n",
        "                siamese_index = np.random.choice(self.label_to_indices[siamese_label])\n",
        "            img2 = self.train_data[siamese_index]\n",
        "        else:\n",
        "            img1 = self.test_data[self.test_pairs[index][0]]\n",
        "            img2 = self.test_data[self.test_pairs[index][1]]\n",
        "            target = self.test_pairs[index][2]\n",
        "\n",
        "        img1 = Image.fromarray(img1.numpy(), mode='L')\n",
        "        img2 = Image.fromarray(img2.numpy(), mode='L')\n",
        "        if self.transform is not None:\n",
        "            img1 = self.transform(img1)\n",
        "            img2 = self.transform(img2)\n",
        "        return (img1, img2), target\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.mnist_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Z1fGqzm1E0M"
      },
      "source": [
        "## 3.3 Loss function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C48BKMc91E0N"
      },
      "source": [
        "$$L_{contrastive}(x_0, x_1, y) = \\frac{1}{2} y \\lVert f(x_0)-f(x_1)\\rVert_2^2 + \\frac{1}{2}(1-y)\\{max(0, m-\\lVert f(x_0)-f(x_1)\\rVert_2)\\}^2$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DCe9G3MkJZSJ"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EpVpPbBZ1E0N"
      },
      "source": [
        "### Exercise\n",
        "\n",
        "Fill in the missing parts of the `contrastive loss`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jiKkNs3e1E0O"
      },
      "source": [
        "class ContrastiveLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Contrastive loss\n",
        "    Takes embeddings of two samples and a target label == 1 if samples are from the same class and label == 0 otherwise\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin):\n",
        "        super(ContrastiveLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "        self.eps = 1e-9\n",
        "\n",
        "    def forward(self, output1, output2, target, size_average=True):\n",
        "        # compute squared distances between output2 and output1  \n",
        "        squared_distances = (output2 - output1).pow(2).sum(1)\n",
        "        # add the second term from them loss. You can use ReLU for compressing the max formula\n",
        "        # losses = 0.5 * (target.float() * squared_distances +\n",
        "        #                  (1-target).float() * F.relu(self.margin - (squared_distances + self.eps).sqrt()).pow(2)) \n",
        "        \n",
        "        losses = 0.5 * (target.float() * squared_distances +\n",
        "                         (1-target).float() * F.relu(self.margin - squared_distances + self.eps))\n",
        "        \n",
        "        return losses.mean() if size_average else losses.sum()\n",
        "        \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7nclgG91E0Q"
      },
      "source": [
        "## 3.4 Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-qpIq-TzCuue"
      },
      "source": [
        "# Set up data loaders\n",
        "siamese_train_dataset = SiameseMNIST(train_dataset) # Returns pairs of images and target same/different\n",
        "siamese_test_dataset = SiameseMNIST(test_dataset)\n",
        "\n",
        "args.batch_size = 128\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True}\n",
        "siamese_train_loader = torch.utils.data.DataLoader(siamese_train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "siamese_test_loader = torch.utils.data.DataLoader(siamese_test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "margin = 1.\n",
        "embedding_net = EmbeddingNet()\n",
        "model = SiameseNet(embedding_net)\n",
        "loss_fn = ContrastiveLoss(margin)\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "\n",
        "args.lr = 1e-3\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fYJbkKAm1E0S"
      },
      "source": [
        "def train_siamese_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=100):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, num_samples =  0, 0\n",
        "  \n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        num_samples += data[0].size(0)\n",
        "        \n",
        "        data = tuple(d.to(args.device) for d in data)\n",
        "        target = target.to(args.device)\n",
        "          \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        outputs = model(data[0], data[1])\n",
        "        # alternatively: outputs = model(*data)\n",
        "        \n",
        "        loss = loss_fn(outputs[0], outputs[1], target)\n",
        "        # alternatively: loss = loss_fn(*outputs, target)\n",
        "        \n",
        "        losses.append(loss.data.item())\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        \n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train: [{}/{} ({:.0f}%)]\\tLoss: {:.6f} '.format(\n",
        "                batch_idx * len(data[0]), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), np.mean(losses)))           \n",
        "            \n",
        "            total_loss += np.sum(losses)\n",
        "            losses = []\n",
        "            \n",
        "    return total_loss/(batch_idx + 1)\n",
        "\n",
        "def test_siamese_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    with torch.no_grad():\n",
        "        model.eval()\n",
        "        losses = []\n",
        "        num_samples = 0\n",
        "    \n",
        "        for batch_idx, (data, target) in enumerate(test_loader):\n",
        "\n",
        "            num_samples += data[0].size(0)\n",
        "            data = tuple(d.to(args.device) for d in data)\n",
        "            target = target.to(args.device)\n",
        "            outputs = model(data[0], data[1])\n",
        "\n",
        "            loss = loss_fn(outputs[0], outputs[1], target)\n",
        "            losses.append(loss.data.item())\n",
        "    \n",
        "        return np.sum(losses)/(batch_idx + 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_IqiBATeCuuh"
      },
      "source": [
        "start_epoch = 0\n",
        "\n",
        "# main training loop\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "\n",
        "    # train stage\n",
        "    train_loss = train_siamese_epoch(siamese_train_loader, model, loss_fn, optimizer, args)\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss)\n",
        "    \n",
        "    # testing/validation stage    \n",
        "    test_loss = test_siamese_epoch(siamese_test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             test_loss)\n",
        "    print(message)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOzgmRO91E0X"
      },
      "source": [
        "## 3.5 Visualizations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rBBBuPjiCuup"
      },
      "source": [
        "train_embeddings_cl, train_labels_cl = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_cl, train_labels_cl, title='Train embeddings (constrastive loss)')\n",
        "test_embeddings_cl, test_labels_cl = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_cl, test_labels_cl, title='Test embeddings (contrastive loss)')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28NPJHe91E0b"
      },
      "source": [
        "In order to two compare vectors $x_1$ and $x_2$ we can use the `cosine similarity` \n",
        "\n",
        "$$\\text{similarity}=\\frac{x_1 \\cdot x_2}{\\max(\\Vert x_1 \\Vert _2 \\cdot \\Vert x_2 \\Vert_2, \\epsilon)}$$\n",
        "\n",
        "An alternative is the Euclidean distance."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ObKaI7Oa1E0b"
      },
      "source": [
        "In order to save computation at query time we can pre-process our vectors and L2-normalize them. Now we can simply perform comparison by dot product"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ncftmsls1E0b"
      },
      "source": [
        "#### Exercise\n",
        "Perform L2-normalization on the embeddings using `numpy` "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rigYIlwM1E0c"
      },
      "source": [
        "# L2-normalize embeddings\n",
        "test_embeddings_norm = ...."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53--KowM64OE"
      },
      "source": [
        "### Question\r\n",
        "Why do we normalize features?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "azGVIKYV1E0e"
      },
      "source": [
        "### Exercise\n",
        "Write now a function `most_sim` that computes all dot products between a query vector and the dataset, extracts the indices of the `topk` most similar vectors and put thme in a list of tuples ("
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HuvS4r6D1E0f"
      },
      "source": [
        "def most_sim(x, emb, topk=6):\n",
        "       return None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hF5DQ3F71E0h"
      },
      "source": [
        "test_images_raw = get_raw_images(test_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TzxTj6jn1E0i"
      },
      "source": [
        "def launch_query(test_embeddings_norm, test_images_raw, query_id=None):\n",
        "    query_id = random.randint(0, test_embeddings_norm.shape[0]) if query_id is None else query_id\n",
        "    query_vector = test_embeddings_norm[query_id,:]\n",
        "\n",
        "    print(f'query_id: {query_id} | query_embedding: {query_vector}')\n",
        "    knns = most_sim(query_vector, test_embeddings_norm)\n",
        "    knn_images = np.array([test_images_raw[x[0]] for x in knns ])\n",
        "    \n",
        "    title=['q: 1.0', f'1nn: {knns[1][1]:.3}', f'2nn: {knns[2][1]:.3}', \n",
        "           f'3nn: {knns[3][1]:.3}', f'4nn: {knns[4][1]:.3}', f'5nn: {knns[5][1]:.3}']\n",
        "    show(torchvision.utils.make_grid(torch.from_numpy(knn_images)), title=title)\n",
        "#     print(knns)\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ydZMwGHt1E0k"
      },
      "source": [
        "for i in range(5):\n",
        "    launch_query(test_embeddings_norm, test_images_raw)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MbKXy6yQCuuu"
      },
      "source": [
        "# Triplet network\n",
        "We'll train a triplet network, that takes an anchor, positive (same class as anchor) and negative (different class than anchor) examples. The objective is to learn embeddings such that the anchor is closer to the positive example than it is to the negative example by some margin value.\n",
        "\n",
        "![alt text](images/anchor_negative_positive.png \"Source: FaceNet\")\n",
        "Source: [2] *Schroff, Florian, Dmitry Kalenichenko, and James Philbin. [Facenet: A unified embedding for face recognition and clustering.](https://arxiv.org/abs/1503.03832) CVPR 2015.*\n",
        "\n",
        "**Triplet loss**:   $L_{triplet}(x_a, x_p, x_n) = max(0, m +  \\lVert f(x_a)-f(x_p)\\rVert_2^2 - \\lVert f(x_a)-f(x_n)\\rVert_2^2$\\)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pBLyLplg1E0m"
      },
      "source": [
        "## 4.1 Architecture\n",
        "We will first define the triplet architecture on top of our `EmbeddingNet`\n",
        "\n",
        "### Exercise\n",
        "\n",
        "Fill in the forward part of `TripleNet`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b2kAbre51E0o"
      },
      "source": [
        "class TripletNet(nn.Module):\n",
        "    def __init__(self, embedding_net):\n",
        "        super(TripletNet, self).__init__()\n",
        "        self.embedding_net = embedding_net\n",
        "\n",
        "    def forward(self, x1, x2, x3):\n",
        "        # missing 3 lines here\n",
        "        \n",
        "        return output1, output2, output3\n",
        "\n",
        "    def get_embedding(self, x):\n",
        "        return self.embedding_net(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yCRCt5Fj1E0q"
      },
      "source": [
        "## 4.2 Data loader\n",
        "We will also need to adapt our data loader to fetch triplets of images "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MvAuRt6J1E0s"
      },
      "source": [
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data.sampler import BatchSampler\n",
        "from PIL import Image\n",
        "\n",
        "class TripletMNIST(Dataset):\n",
        "    \"\"\"\n",
        "    Train: For each sample (anchor) randomly chooses a positive and negative samples\n",
        "    Test: Creates fixed triplets for testing\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, mnist_dataset):\n",
        "        self.mnist_dataset = mnist_dataset\n",
        "        self.train = self.mnist_dataset.train\n",
        "        self.transform = self.mnist_dataset.transform\n",
        "\n",
        "        if self.train:\n",
        "            self.train_labels = self.mnist_dataset.train_labels\n",
        "            self.train_data = self.mnist_dataset.train_data\n",
        "            self.labels_set = set(self.train_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.train_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "        else:\n",
        "            self.test_labels = self.mnist_dataset.test_labels\n",
        "            self.test_data = self.mnist_dataset.test_data\n",
        "            # generate fixed triplets for testing\n",
        "            self.labels_set = set(self.test_labels.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.test_labels.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "            random_state = np.random.RandomState(29)\n",
        "\n",
        "            triplets = [[i,\n",
        "                         random_state.choice(self.label_to_indices[self.test_labels[i].item()]),\n",
        "                         random_state.choice(self.label_to_indices[\n",
        "                                                 np.random.choice(\n",
        "                                                     list(self.labels_set - set([self.test_labels[i].item()]))\n",
        "                                                 )\n",
        "                                             ])\n",
        "                         ]\n",
        "                        for i in range(len(self.test_data))]\n",
        "            self.test_triplets = triplets\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        if self.train:\n",
        "            img1, label1 = self.train_data[index], self.train_labels[index].item()\n",
        "            positive_index = index\n",
        "            while positive_index == index:\n",
        "                positive_index = np.random.choice(self.label_to_indices[label1])\n",
        "            negative_label = np.random.choice(list(self.labels_set - set([label1])))\n",
        "            negative_index = np.random.choice(self.label_to_indices[negative_label])\n",
        "            img2 = self.train_data[positive_index]\n",
        "            img3 = self.train_data[negative_index]\n",
        "        else:\n",
        "            img1 = self.test_data[self.test_triplets[index][0]]\n",
        "            img2 = self.test_data[self.test_triplets[index][1]]\n",
        "            img3 = self.test_data[self.test_triplets[index][2]]\n",
        "\n",
        "        img1 = Image.fromarray(img1.numpy(), mode='L')\n",
        "        img2 = Image.fromarray(img2.numpy(), mode='L')\n",
        "        img3 = Image.fromarray(img3.numpy(), mode='L')\n",
        "        if self.transform is not None:\n",
        "            img1 = self.transform(img1)\n",
        "            img2 = self.transform(img2)\n",
        "            img3 = self.transform(img3)\n",
        "        return (img1, img2, img3), []\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.mnist_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "779Hh05A1E0u"
      },
      "source": [
        "## 4.3 Loss function\n",
        "\n",
        "### Exercise\n",
        "\n",
        "Fill in the missing parts of the `triplet loss`:\n",
        " $L_{triplet}(x_a, x_p, x_n) = max(0, m +  \\lVert f(x_a)-f(x_p)\\rVert_2^2 - \\lVert f(x_a)-f(x_n)\\rVert_2^2$\\)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kRaxvIFJ1E0u"
      },
      "source": [
        "class TripletLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Triplet loss\n",
        "    Takes embeddings of an anchor sample, a positive sample and a negative sample\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin):\n",
        "        super(TripletLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "\n",
        "    def forward(self, anchor, positive, negative, size_average=True):\n",
        "        distance_positive = None # fill in code\n",
        "        distance_negative = None  # fill in code\n",
        "        # you can again use ReLU instead of max         \n",
        "        losses = None # fill in code\n",
        "        return losses.mean() if size_average else losses.sum()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PyCDVrG81E0w"
      },
      "source": [
        "## 4.4 Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jv4DvFucCuuu"
      },
      "source": [
        "triplet_train_dataset = TripletMNIST(train_dataset) # Returns triplets of images\n",
        "triplet_test_dataset = TripletMNIST(test_dataset)\n",
        "\n",
        "args.batch_size = 128\n",
        "kwargs = {'num_workers': args.num_workers, 'pin_memory': True} \n",
        "triplet_train_loader = torch.utils.data.DataLoader(triplet_train_dataset, batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "triplet_test_loader = torch.utils.data.DataLoader(triplet_test_dataset, batch_size=args.batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "margin = 1.\n",
        "embedding_net = EmbeddingNet()\n",
        "model = TripletNet(embedding_net)\n",
        "loss_fn = TripletLoss(margin)\n",
        "model.to(args.device)\n",
        "loss_fn.to(args.device)\n",
        "args.lr = 1e-3\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr)\n",
        "scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)\n",
        "n_epochs = 5\n",
        "log_interval = 100"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AsMAV5Kt1E0z"
      },
      "source": [
        "### Exercise\n",
        "\n",
        "Code your own train/test sequences similarly to the previous examples.\n",
        "Watch out for some differences though."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HIeuIH4y1E00"
      },
      "source": [
        "def train_triplet_epoch(train_loader, model, loss_fn, optimizer, args, log_interval=100):\n",
        "    model.train()\n",
        "    losses = []\n",
        "    total_loss, num_samples =  0, 0\n",
        "  \n",
        "    # fill in code here\n",
        "    \n",
        "    return total_loss/(batch_idx + 1)\n",
        "\n",
        "def test_triplet_epoch(test_loader, model, loss_fn, args, log_interval=50):\n",
        "    losses = []\n",
        "    num_samples = 0\n",
        "    # fill in code here\n",
        "\n",
        "    return np.sum(losses)/(batch_idx + 1)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6TfLZH1V1E02"
      },
      "source": [
        "start_epoch = 0\n",
        "\n",
        "\n",
        "# main training loop\n",
        "for epoch in range(start_epoch, args.num_epochs):\n",
        "\n",
        "    # train stage\n",
        "    train_loss = train_triplet_epoch(triplet_train_loader, model, loss_fn, optimizer, args)\n",
        "    message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}'.format(\n",
        "        epoch + 1, args.num_epochs, train_loss)\n",
        "    \n",
        "    # testing/validation stage    \n",
        "    test_loss = test_triplet_epoch(triplet_test_loader, model, loss_fn, args)\n",
        "    \n",
        "    message += '\\nEpoch: {}/{}. Validation set: Average loss: {:.4f}'.format(epoch + 1, args.num_epochs,\n",
        "                                                                             test_loss)\n",
        "    print(message)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LEkIWNDq1E05"
      },
      "source": [
        "## 4.5 Visualizations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ysh4Ry7ZCuu_"
      },
      "source": [
        "train_embeddings_tl, train_labels_tl = extract_embeddings(train_loader, model, args)\n",
        "plot_embeddings(train_embeddings_tl, train_labels_tl, title='Train triplet embeddings')\n",
        "test_embeddings_tl, test_labels_tl = extract_embeddings(test_loader, model, args)\n",
        "plot_embeddings(test_embeddings_tl, test_labels_tl, title='Val triplet embeddings')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zw0y8bgB1E08"
      },
      "source": [
        "# L2-normalize embeddings\n",
        "test_embeddings_tl_norm = test_embeddings_tl / np.linalg.norm(test_embeddings_tl, axis=-1, keepdims=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YnNl2ffL1E09"
      },
      "source": [
        "test_images_raw = get_raw_images(test_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h9JOhZTY1E0_"
      },
      "source": [
        "for i in range(5):\n",
        "    launch_query(test_embeddings_tl_norm, test_images_raw)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}