{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Можем ли мы использовать классификацию изображений для решения задачи поиска похожих изображений? И возможно ли обучить для этой задачи довольно надежную, глубокую CNN? Причем ее результат должен быть стабильным при разных ракурсах, освещении, позе, и.т.д.\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/metric-learning/img1.png)\n",
    "\n",
    "Обычно это достигается, когда у нас есть хороший объем данных. Например, $1000$ примеров для каждого класса. В классификации изображений, если количество точек данных в классе невелико, это может привести к переобучению и очень плохим результатам. Кроме того, классификация изображений обычно хорошо работает, когда количество классов невелико.\n",
    "\n",
    "Однако, как правило, это не относится к задачам проверки личности / изображения. Тут все совсем наоборот. Здесь у нас обычно очень большое количество классов, а количество примеров на класс довольно мало. И здесь на помощь приходит few-shot/one-shot learning.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Few-shot Learning\n",
    "Это задача классификации, цель которой - научиться отличать классы объектов по одному / нескольким обучающим примерам. Например, в задаче распознавания лиц по одному примеру человека нужно научиться его узнать.\n",
    "\n",
    "Модели глубокого обучения не будут работать с одним учебным примером. Как решить эту проблему? Вместо решения исходной задачи будем обучать **функцию \"похожести\"** разных примеров. То есть **метрику** для их сравнения.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Простая модель\n",
    "Сначала попробуем обучить простую модель для классификации людей по лицам. И будем использовать выходы предпоследнего слоя для обучения метода ближайших соседей и нахождения похожих людей."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "import albumentations as A\n",
    "import albumentations.pytorch\n",
    "import lightning as L\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchmetrics\n",
    "import torchvision\n",
    "from PIL import Image\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from torch.utils import data\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "# Increase these if figures appear small\n",
    "plt.rcParams[\"figure.figsize\"] = fx, fy = (14.08, 6.40)\n",
    "\n",
    "# Not using `bfloat16` matrix multiplication for consistency\n",
    "# You might get better performance without much precision loss\n",
    "# by setting this to \"medium\" on some devices\n",
    "torch.set_float32_matmul_precision(\"high\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    DEVICE = torch.device(\"cuda\")\n",
    "    print(\"Using the GPU 😊\")\n",
    "else:\n",
    "    DEVICE = torch.device(\"cpu\")\n",
    "    print(\"Using the CPU 😞\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(\"lfw\"):\n",
    "    !curl -O https://courses.cv-gml.ru/storage/seminars/metric-learning/lfw.tgz\n",
    "    !tar -xf lfw.tgz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_EPOCHS = 5\n",
    "BATCH_SIZE = 64\n",
    "EMB_DIM = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Напишем датапровайдер. При этом в нем будет функция `make_slice`, которая в зависимости от параметра `leave` либо только оставляет, либо наоборот убирает индексы из списка `base_elems`.\n",
    "\n",
    "Мы будем использовать набор даных [`LFW`](http://vis-www.cs.umass.edu/lfw/). Сам по себе этот набор, как правило, используется только как бенчмарк, а не для обучения полноценных моделей. Однако реальные датасеты для обучения (например, [`CelebA`](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) имеют очень большой размер и они слишком большие для данного семинара. Однако из-за того, что мы используем плохой датасет (и плохую модель), получаемое качество моделей будет не очень хорошим."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_by_index(elems, index):\n",
    "    targets = [], []\n",
    "    for i, elem in enumerate(elems):\n",
    "        targets[i in index].append(elem)\n",
    "    return targets\n",
    "\n",
    "\n",
    "class LFWDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, folder, people_list, transforms):\n",
    "        super().__init__()\n",
    "        self.people_list = people_list\n",
    "        self.need_classes = [i for i in range(len(people_list))]\n",
    "        self.images = []\n",
    "        self.labels = []\n",
    "        self.weights = []\n",
    "        self.class_to_idxs = {}\n",
    "        for i, person in enumerate(people_list):\n",
    "            self.class_to_idxs[i] = []\n",
    "            fnames = sorted(os.listdir(os.path.join(folder, person)))\n",
    "            for fname in fnames:\n",
    "                image = Image.open(os.path.join(folder, person, fname))\n",
    "                image = np.asarray(image.convert(\"RGB\"))\n",
    "                self.images += [image]\n",
    "                self.labels += [i]\n",
    "                self.weights += [1.0 / len(fnames)]\n",
    "                self.class_to_idxs[i].append(len(self.images) - 1)\n",
    "        self.transform = transforms\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        target = self.labels[index]\n",
    "        sample = self.transform(image=self.images[index])[\"image\"]\n",
    "        return sample, target\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def make_slice(self, base_elems, leave=True):\n",
    "        self.images = split_by_index(self.images, base_elems)[leave]\n",
    "        self.weights = split_by_index(self.weights, base_elems)[leave]\n",
    "        self.labels = split_by_index(self.labels, base_elems)[leave]\n",
    "\n",
    "        self.class_to_idxs = {i: [] for i in range(len(self.people_list))}\n",
    "        for idx, label in enumerate(self.labels):\n",
    "            self.class_to_idxs[label].append(idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Построим списки людей для обучения и тестирования."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_people = []\n",
    "test_people = []\n",
    "for person in sorted(os.listdir(\"lfw\")):\n",
    "    cnt = len(os.listdir(\"lfw/\" + person))\n",
    "    if cnt >= 3:\n",
    "        train_people.append(person)\n",
    "    elif cnt >= 2:\n",
    "        test_people.append(person)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим аументацию данных. Сделаем три набора данных:\n",
    "- train_data - обучающая выборка - все картинки, встречающиеся хотя бы 3 раза\n",
    "- test_data_index - фото по одному кадру людей из тестовой выборки (не лежащие в ней) (одно из двух фото)\n",
    "- test_data - фото кадров людей из тестовой выборки (одно из двух фото)\n",
    "\n",
    "Построим для этих наборов `DataLoader`-ы."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_transform = A.Compose(\n",
    "    [\n",
    "        A.HorizontalFlip(0.5),\n",
    "        A.Rotate(15),\n",
    "        A.RandomResizedCrop((224, 224), scale=(0.6, 1.0)),\n",
    "        A.RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=0.5),\n",
    "        A.RandomBrightnessContrast(0.1, 0.1, p=0.5),\n",
    "        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n",
    "        A.pytorch.ToTensorV2(),\n",
    "    ],\n",
    ")\n",
    "\n",
    "test_transform = A.Compose(\n",
    "    [\n",
    "        A.Resize(224, 224),\n",
    "        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n",
    "        A.pytorch.ToTensorV2(),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = LFWDataset(\"lfw\", train_people, train_transform)\n",
    "test_data_index = LFWDataset(\"lfw\", test_people, test_transform)\n",
    "test_data = LFWDataset(\"lfw\", test_people, test_transform)\n",
    "\n",
    "random.seed(51)\n",
    "base_elems = []\n",
    "for clas, idxs in test_data_index.class_to_idxs.items():\n",
    "    idx = random.choice(idxs)\n",
    "    base_elems.append(idx)\n",
    "test_data_index.make_slice(base_elems, True)\n",
    "test_data.make_slice(base_elems, False)\n",
    "\n",
    "sampler = data.WeightedRandomSampler(\n",
    "    train_data.weights,\n",
    "    len(train_data.weights),\n",
    ")\n",
    "train_dl = data.DataLoader(\n",
    "    train_data,\n",
    "    BATCH_SIZE,\n",
    "    sampler=sampler,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "oneshot_dl = data.DataLoader(\n",
    "    test_data_index,\n",
    "    BATCH_SIZE,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "test_dl = data.DataLoader(\n",
    "    test_data,\n",
    "    BATCH_SIZE,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Визуализируем получающиеся у нас данные."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_batch(dl):\n",
    "    for batch in dl:\n",
    "        images, labels = batch\n",
    "        fig, ax = plt.subplots(figsize=(7.5, 7.5))\n",
    "        ax.set_yticks([])\n",
    "        ax.set_xticks([])\n",
    "        sample = images[:20]\n",
    "        sample = sample - sample.min()\n",
    "        sample = sample / sample.max()\n",
    "        ax.imshow(make_grid(sample, nrow=5).permute(1, 2, 0))\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_batch(train_dl)\n",
    "show_batch(test_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseEmbeddingTrainer(L.LightningModule):\n",
    "    def __init__(self, *, emb_net, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.emb_net = emb_net\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=1e-3)\n",
    "\n",
    "    def _log_metrics(self, loss, accs, kind):\n",
    "        metrics = {}\n",
    "        if loss is not None:\n",
    "            metrics[f\"{kind}_loss\"] = loss\n",
    "        if accs is not None:\n",
    "            metrics[f\"{kind}_accs\"] = accs\n",
    "        self.log_dict(\n",
    "            metrics,\n",
    "            prog_bar=True,\n",
    "            logger=True,\n",
    "            on_step=kind == \"train\",\n",
    "            on_epoch=True,\n",
    "        )\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим модель `EmbeddingNet` для извлечения признаков изображений людей."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingNet(nn.Sequential):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1\n",
    "        self.convnet = torchvision.models.resnet18(weights=weights)\n",
    "\n",
    "        linear_size = self.convnet.fc.in_features\n",
    "        self.convnet.fc = nn.Linear(linear_size, EMB_DIM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим модель `ClassificationTrainer` для построения классификатора людей где на каждого человека отведен свой класс."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClassificationTrainer(BaseEmbeddingTrainer):\n",
    "    def __init__(self, *, train_classes, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.accuracy = torchmetrics.Accuracy(\n",
    "            task=\"multiclass\",\n",
    "            num_classes=train_classes,\n",
    "        )\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        self.classifier = nn.Linear(\n",
    "            EMB_DIM,\n",
    "            train_classes,\n",
    "        )\n",
    "\n",
    "    def training_step(self, batch):\n",
    "        x, y = batch\n",
    "\n",
    "        f = self.emb_net(x)\n",
    "        p = self.classifier(f)\n",
    "\n",
    "        loss = F.cross_entropy(p, y)\n",
    "        accs = self.accuracy(p.argmax(dim=-1), y)\n",
    "\n",
    "        return self._log_metrics(loss, accs, \"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обучим модель"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(\n",
    "    model,\n",
    "    experiment_path,\n",
    "    train_dl,\n",
    "    max_epochs=MAX_EPOCHS,\n",
    "    **trainer_kwargs,\n",
    "):\n",
    "    callbacks = [\n",
    "        L.pytorch.callbacks.TQDMProgressBar(leave=True),\n",
    "        L.pytorch.callbacks.LearningRateMonitor(),\n",
    "    ]\n",
    "    trainer = L.Trainer(\n",
    "        callbacks=callbacks,\n",
    "        max_epochs=max_epochs,\n",
    "        default_root_dir=experiment_path,\n",
    "        **trainer_kwargs,\n",
    "    )\n",
    "    trainer.fit(model, train_dl)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = train_model(\n",
    "    ClassificationTrainer(\n",
    "        emb_net=EmbeddingNet(),\n",
    "        train_classes=len(train_data.people_list),\n",
    "    ),\n",
    "    \"runs/train_cls\",\n",
    "    train_dl,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обучим k-NN на изображениях из идекса людей. Протестируем ее на наборе из тех же людей, но других их изображениях."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_emb_net(emb_net, dataloader, normalize=False):\n",
    "    ### YOUR CODE HERE\n",
    "    data_x = []\n",
    "    data_y = []\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in dataloader:\n",
    "            feats = emb_net(inputs.to(DEVICE))\n",
    "            if normalize:\n",
    "                feats = F.normalize(feats)\n",
    "            feats = feats.detach().cpu().numpy()\n",
    "            labels = labels.detach().cpu().numpy()\n",
    "            data_x.append(feats)\n",
    "            data_y.append(labels)\n",
    "\n",
    "    data_x = np.concatenate(data_x, axis=0)\n",
    "    data_y = np.concatenate(data_y, axis=0)\n",
    "    return data_x, data_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_knn(emb_net, normalize=False):\n",
    "    ### YOUR CODE HERE\n",
    "    data_x, data_y = run_emb_net(emb_net, oneshot_dl, normalize)\n",
    "    knn = KNeighborsClassifier(n_neighbors=1)\n",
    "    knn = knn.fit(data_x, data_y)\n",
    "    return knn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_knn(emb_net, knn, normalize=False):\n",
    "    ### YOUR CODE HERE\n",
    "    data_x, data_y = run_emb_net(emb_net, test_dl, normalize)\n",
    "\n",
    "    total_acc = 0\n",
    "    total_cnt = 0\n",
    "    for feat, label in zip(data_x, data_y):\n",
    "        pred = knn.predict(feat[None]).squeeze(0)\n",
    "        total_acc += pred == label\n",
    "        total_cnt += 1\n",
    "\n",
    "    acc = total_acc / total_cnt\n",
    "    print(f\"Accuracy = {acc:.2%} ({total_acc} / {total_cnt})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.eval()\n",
    "net.emb_net.to(DEVICE)\n",
    "knn = train_knn(net.emb_net)\n",
    "test_knn(net.emb_net, knn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Метрика:\n",
    "Метрика - это неотрицательная функция $d (x, y)$ между двумя точками $x$ и $y$, которая описывает так называемое понятие \"расстояния\" между этими двумя точками. Формально, есть несколько свойств, которым должна удовлетворять метрика:\n",
    "- Неотрицательность: $d (x, y) ≥ 0$ и $d (x, y) = 0$, если $x = y$.\n",
    "- Ннеравенство треугоьлника: $d (x, y) ≤ d (x, z) + d (z, y)$.\n",
    "- Симметричность: $g (x, y) = g (y, x)$.\n",
    "\n",
    "### Примеры:\n",
    "- Евклидова метрика: в d-мерном векторном пространстве\n",
    "$$d(x, y) = \\sqrt{\\sum_{i=1}^d (x_i-y_i)^2}$$\n",
    "- Дискретная метрика:\n",
    "$$\\begin{cases}\n",
    "    1, x \\neq y\\\\\n",
    "    0, x == y\n",
    "\\end{cases} \n",
    "$$\n",
    "\n",
    "### Обучение метрик:\n",
    "\n",
    "Что делает базовый алгоритм машинного обучения? - Учитывая данные и соответствующие выходные метки, настраивается набор правил или некоторая сложная функция, которая входным данным сопоставляет соответствующие выходные метки.\n",
    "\n",
    "Одним из простейших алгоритмов машинного обучения, в котором фиксируется информация о расстоянии, является алгоритм $kNN$, идея которого состоит в том, чтобы найти окрестность $k$ ближайших точек данных для новой точки. Далее этой точке назначается класс, которому принадлежит большинство из $k$ точек окрестности.\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/metric-learning/img2.png)\n",
    "\n",
    "У метрического обучения похожая цель - изучить функцию похожести на основе данных. Метрическое обучение направлено на изучение встраиваемых данных / векторов признаков таким образом, чтобы уменьшить расстояние между векторами признаков, соответствующими примерам, принадлежащим одному классу, и увеличить расстояние между векторами признаков, соответствующими разным классам.\n",
    "\n",
    "Евклидовы расстояния менее значимы в больших измерениях. Им не удается уловить нелинейность данных по ряду причин:\n",
    "- Они представляют собой изотропную (одинаковую во всех направлениях) метрику расстояния.\n",
    "- Эти расстояния не отражают структуру классов.\n",
    "\n",
    "Можно использовать неизотропные расстояния, которые используют свойства данных для определения структуры отношений классов. Расстояние Махаланобиса - одна из таких метрик:\n",
    "\n",
    "$$d(x, y) = \\sqrt{(x-y)^TM(x-y)}$$\n",
    "\n",
    "Здесь $M$ является обратной матрицей ковариации и действует как весовой член квадрата евклидова расстояния. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deep metric learning\n",
    "Deep metric learning использует нейронные сети для автоматического изучения отличительных характеристик изображений и вычисления метрики их похожести.\n",
    "\n",
    "Зачем нужно глубокое метрическое обучение?\n",
    "- Нужна устойчивость моделей при смене позы/ракурса/освещения.\n",
    "- Это прикладная задача. Например, создание эффективной системы проверки лиц. При этом не хочется обучать её полностью с нуля с каждым новым примером в системе.\n",
    "- Классическое Metric Learning имеет лишь ограниченную возможность фиксировать нелинейность данных.\n",
    "- Глубокое метрическое обучение помогает фиксировать нелинейную структуру данных.\n",
    "\n",
    "Далее рассмотрим функции потерь для обучения таких нейронных сетей.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Siamese Networks:\n",
    "\n",
    "Одна из фундаментальных идей - это сиамская нейросетевая модель. Сиамская сеть - это архитектура симметричной нейронной сети, состоящая из двух идентичных подсетей, которые используют одни и те же наборы параметров и обучаются, вычисляя метрику между по последним слоям каждой подсети.\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/metric-learning/img3.png)\n",
    "\n",
    "Обозначим преобразование, выполняемое нейросетью, как $f(x)$, а расстояние между примерами $d (f(x), f(y))$. Наша цель состоит в следующем:\n",
    "\n",
    "- Если $x$ и $y$ - примеры одного и того же класса, то $d (f(x), f(y))$ мало,\n",
    "- Если $x$ и $y$ - примеры разных классов $d (f(x), f(y))$ большое.\n",
    "\n",
    "Сиамские сети часто используются в сочетании с **Contrastive loss**.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим класс `SiameseTrainer` для обучения сиамских сетей."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiameseTrainer(BaseEmbeddingTrainer):\n",
    "    def __init__(self, *, loss_functor, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.loss_functor = loss_functor\n",
    "\n",
    "    def training_step(self, batch):\n",
    "        xs, ys = batch\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        fs = [self.emb_net(x) for x in xs]\n",
    "        loss = self.loss_functor(fs, ys)\n",
    "\n",
    "        return self._log_metrics(loss, None, \"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Contrastive loss:\n",
    "Эта loss-функция обучает вектора признаков таким образом, чтобы расстояние между двумя примерами из одного и того же класса было небольшим, а расстояние между разными классами было большим. А теперь давайте разберемся с математикой:\n",
    "$$L=L(f(x), f(y), z) = z ||f(x) - f(y)||_2^2 + (1-z)max\\{0, m - ||f(x) - f(y)||_2\\}^2$$\n",
    "\n",
    "Здесь $m> 0$ - отступ, $z$ - логическое значение, равное 1, если $x$, $y$ - примеры одного класса, и $0$ в противном случае, $||*||_2$ - норма $L_2$.\n",
    "\n",
    "Давайте разберемся в этом уравнении. Если $x$ и $y$ являются парами одного класса и $z$ равно $1$, то второй член исчезает ($m$ не влияет). И функция минимизирует первое слагаемое - расстояние между примерами одного класса.\n",
    "\n",
    "Теперь, если $x$ и $y$ - это пары разного класса и $z$ равно $0$, то первый член равен нулю. Если сеть считает, что $x$ и $y$ похожи, а расстояние между ними невелико, то мы получим значение loss-функции, близкое к $m$.\n",
    "\n",
    "Роль отступа $m$: Следует отметить, пары разных классов дадут вклад в функцию потерь только в том случае, если расстояние $||f(x) - f(y)||_2^2 < m$. Это означает, что функцию не будет заботить, как далеко отойдут отрицательные пары $x$ и $y$, когда этот предел $m$ достигнет. Таким образом, он может больше сосредоточиться на трудных точках.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим класс `ContrastiveLFW`, который будет предоставлять данные для обучения сиамской сети. Он будет для данного изображения дополнительно генерировать с вероятностью `0.5` либо другой пример того же класса, либо пример другого класса."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_choice_except(options, exception):\n",
    "    ### YOUR CODE HERE\n",
    "    choice = exception\n",
    "    while choice == exception:\n",
    "        choice = random.choice(options)\n",
    "    return choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContrastiveLFW(torch.utils.data.Dataset):\n",
    "    def __init__(self, lfw_dataset):\n",
    "        super().__init__()\n",
    "        self.lfw_dataset = lfw_dataset\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        ### YOUR CODE HERE\n",
    "        sample1, target1 = self.lfw_dataset[index]\n",
    "\n",
    "        if random.random() > 0.5:\n",
    "            other_target = target1\n",
    "        else:\n",
    "            other_target = random_choice_except(\n",
    "                self.lfw_dataset.need_classes,\n",
    "                exception=target1,\n",
    "            )\n",
    "\n",
    "        other_index = random_choice_except(\n",
    "            self.lfw_dataset.class_to_idxs[other_target],\n",
    "            exception=index,\n",
    "        )\n",
    "        sample2, target2 = self.lfw_dataset[other_index]\n",
    "        return [sample1, sample2], [target1, target2]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.lfw_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим провайдеры данных."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_siamese_data = ContrastiveLFW(train_data)\n",
    "contr_train_dl = data.DataLoader(\n",
    "    train_siamese_data,\n",
    "    BATCH_SIZE,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Реализуем loss-функцию."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContrastiveLoss(nn.Module):\n",
    "    def __init__(self, margin):\n",
    "        super().__init__()\n",
    "        self.margin = margin\n",
    "        self.eps = 1e-9\n",
    "\n",
    "    def forward(self, fs, ys):\n",
    "        f1, f2 = fs\n",
    "        y1, y2 = ys\n",
    "\n",
    "        m = self.margin\n",
    "        eps = self.eps\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        dist_pos = (f2 - f1).square().sum(axis=-1)\n",
    "        dist_neg = F.relu(m - (dist_pos + eps).sqrt()).square()\n",
    "        loss = 0.5 * torch.where(y1 == y2, dist_pos, dist_neg)\n",
    "        return loss.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обучим сеть."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_contr = train_model(\n",
    "    SiameseTrainer(\n",
    "        emb_net=EmbeddingNet(),\n",
    "        loss_functor=ContrastiveLoss(margin=2.0),\n",
    "    ),\n",
    "    \"runs/siamese_contrastive\",\n",
    "    contr_train_dl,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Протестируем качество k-NN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_contr.eval()\n",
    "net_contr.emb_net.to(DEVICE)\n",
    "knn = train_knn(net_contr.emb_net)\n",
    "test_knn(net_contr.emb_net, knn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Triplet loss\n",
    "\n",
    "Триплетная сеть состоит из трех идентичных подсетей, которые имеют одинаковые наборы параметров. Обучение выполняется в виде набора из трех изображений: 1) Якорное изображение (базовое изображение), 2) Позитивное изображение (принадлежащее к тому же классу, что и якорь), 3) Негативное изображение (изображение, не принадлежащее классу якоря).\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/metric-learning/img4.png)\n",
    "\n",
    "Цель triplet loss - гарантировать, что якорное изображение близко ко всем изображениям того же класса. Отрицательных примеры должны быть дальше, чем положительные.\n",
    "\n",
    "$$L=L(f(A), f(B), f(C)) = max\\{0, ||f(A) - f(B)||_2^2 - ||f(A) - f(C)||_2^2 + m\\}$$\n",
    "В этой формуле $A$ - якорь, $B$ - позитивный пример, $C$ - негативный пример, $m$ - параметр, $||*||_2$ - норма $L_2$.\n",
    "\n",
    "Давайте разберемся в этом уравнении. Когда loss-функция будет равна $0$? Второй член внутри $max(*)$ должен быть $<=0$, то есть $d(A, N) >= d(A, B) + m$. Если это произойдет, то на этой тройке параметры нейросети при обучении не обновятся.\n",
    "\n",
    "Когда loss-функция будет положительна? Когда, $d (A, N) < d (A, B) + m$. Полезными и трудными для обучения являются тройки, в которых $d (A, N) << d (A, B)$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим класс `TripletLFW`, который будет предоставлять данные для обучения сиамской сети. Он будет для данного изображения дополнительно генерировать положительный пример того же класса и отрицательный пример другого класса."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TripletLFW(torch.utils.data.Dataset):\n",
    "    def __init__(self, lfw_dataset):\n",
    "        super().__init__()\n",
    "        self.lfw_dataset = lfw_dataset\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        ### YOUR CODE HERE\n",
    "        sample1, target1 = self.lfw_dataset[index]\n",
    "\n",
    "        positive_index = random_choice_except(\n",
    "            self.lfw_dataset.class_to_idxs[target1],\n",
    "            exception=index,\n",
    "        )\n",
    "        sample2, target2 = self.lfw_dataset[positive_index]\n",
    "\n",
    "        negative_target = random_choice_except(\n",
    "            self.lfw_dataset.need_classes,\n",
    "            exception=target1,\n",
    "        )\n",
    "        negative_index = random.choice(\n",
    "            self.lfw_dataset.class_to_idxs[negative_target],\n",
    "        )\n",
    "        sample3, target3 = self.lfw_dataset[negative_index]\n",
    "\n",
    "        return [sample1, sample2, sample3], [target1, target2, target3]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.lfw_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим провайдеры данных."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_triplet_data = TripletLFW(train_data)\n",
    "triplet_train_dl = data.DataLoader(\n",
    "    train_triplet_data,\n",
    "    BATCH_SIZE,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Реализуем loss-функцию."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TripletLoss(nn.Module):\n",
    "    def __init__(self, margin):\n",
    "        super().__init__()\n",
    "        self.margin = margin\n",
    "\n",
    "    def forward(self, fs, ys):\n",
    "        del ys\n",
    "        anchor, positive, negative = fs\n",
    "        m = self.margin\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        dist_pos = (anchor - positive).square().sum(axis=-1)\n",
    "        dist_neg = (anchor - negative).square().sum(axis=-1)\n",
    "        loss = F.relu(dist_pos - dist_neg + m)\n",
    "        return loss.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обучим сеть."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_triplet = train_model(\n",
    "    SiameseTrainer(\n",
    "        emb_net=EmbeddingNet(),\n",
    "        loss_functor=TripletLoss(margin=2.0),\n",
    "    ),\n",
    "    \"runs/siamese_triplet\",\n",
    "    triplet_train_dl,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Протестируем качество k-NN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_triplet.eval()\n",
    "net_triplet.emb_net.to(DEVICE)\n",
    "knn = train_knn(net_triplet.emb_net)\n",
    "test_knn(net_triplet.emb_net, knn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ArcFace loss\n",
    "Данная функция потерь для обучения метрик основана на другой идее. Тут предлагается сравнивать расстояния не по $L_2$ метрике, а по углу до вектора, представляющего нужный класс.\n",
    "\n",
    "Обычная `Softmax`-функция потерь имеет вид:\n",
    "$$L = -\\frac{1}{M}\\sum_{i=1}^M log(\\frac{e^{W^T_{y_i} f_i}}{\\sum_{k=1}^C e^{W^T_{k} f_i}})$$\n",
    "Где $f_i$ - это вектор-признаки, подающиеся на последний полносвязный слой, а $W$ - матрица весов линейного преобразования этого слоя.\n",
    "\n",
    "Рассмотрим задачу классификации на два класса, обученную с использованием функции потерь softmax. Пусть $w_1$ и $w_2$ - векторы для классификации на классы 1 и 2 соответственно. Предположим, что **bias=0**. Точка данных $x_i$ прогнозируется как класс 1, если для вектора признаков $f_i$ выполнено (ниже скалярные произведения):\n",
    "$$||w_1|| * ||f_i|| * cos(\\theta_1) > ||w_2|| * ||f_i|| * cos(\\theta_2)$$\n",
    "Предположим, что $||w_1||=||w_2||=1$, тогда выражение сверху эквивалентно:\n",
    "$$cos(\\theta_1) > cos(\\theta_2)$$\n",
    "То есть на самом деле нам интересен только угол до ближайшего вектора, представляющего класс.\n",
    "\n",
    "![](https://courses.cv-gml.ru/storage/seminars/metric-learning/img5.png)\n",
    "\n",
    "Чтобы правильно обучить такие признаки, норма $W$ обязательно должна быть одинакова для всех классов, $||w_i|| = 1$, $bias=0$. На этапе обучения мы фиксируем $||f|| = s$ ($s$ называется \"параметр масштабирования\"). Так удаляются вариации в радиальных направлениях и разделение происходит только в угловом пространстве на гиперсфере с радиусом s.\n",
    "\n",
    "На этапе тестирования оценка схожести пары лиц рассчитывается в соответствии с косинусным сходством между двумя векторами признаков. При этом косинус угла между векторами $x$ и $y$ равен $\\theta_{x,y}=\\frac{x^Ty}{||x||*||y||}$, а евклидово расстояние $||x-y|| = sqrt(x^Tx+y^Ty-2*x^Ty)$. Если положит $||x||=||y||=1$, то $\\theta_{x,y}=x^Ty$, а $||x-y|| = sqrt(2-2*x^Ty)$.  Когда косинус растет, Евклидово расстояние уменьшается, и наоборот. Поэтому можно продолжать использовать евклидово расстояние в $k-NN$, только надо не забыть нормировать векторы.\n",
    "\n",
    "Основанная на этих идеях функция потерь *ArcFace loss* имеет вид:\n",
    "\n",
    "$$L_{arcface} = -\\frac{1}{M}\\sum_{i=1}^M log(\\frac{e^{s*cos(\\theta_{y_i, i}+m)}}{e^{s*cos(\\theta_{y_i, i}+m)} + \\sum_{k \\neq y_i}^C e^{s*cos(\\theta_{k, i})}})$$\n",
    "\n",
    "Она минимизируется, когда $cos(\\theta_{y_i, i}+m) > cos(\\theta_{k, i})$ при $k \\neq y_i$. Это означает, что $\\theta_{y_i, i}+m < \\theta_{k, i}$, то есть угол до $y_i$-го класса меньше, чем до другого $k$-го с учетом отступа $m$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Реализуем loss-функцию и модель."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArcLoss(nn.Module):\n",
    "    def __init__(self, emb_dim, num_centroids):\n",
    "        super().__init__()\n",
    "        self.scale = 30.0\n",
    "        self.margin = 0.4\n",
    "        self.eps = 1e-7\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        self.centroids = nn.Parameter(torch.empty((num_centroids, emb_dim)))\n",
    "        nn.init.kaiming_uniform_(self.centroids, a=np.sqrt(5))\n",
    "\n",
    "    def forward(self, f, y):\n",
    "        w = self.centroids\n",
    "        s = self.scale\n",
    "        m = self.margin\n",
    "        eps = self.eps\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        w = F.normalize(w, p=2, dim=1)\n",
    "        f = F.normalize(f, p=2, dim=1)\n",
    "        p = F.linear(f, w)\n",
    "        # p.shape == (batch_size, num_centroids)\n",
    "\n",
    "        positives = p.gather(dim=-1, index=y[..., None])\n",
    "        positives = positives.clamp(-1 + eps, 1 - eps)\n",
    "        positives = (positives.acos() + m).cos()\n",
    "\n",
    "        with_margin = p.scatter(dim=-1, index=y[..., None], src=positives)\n",
    "        with_margin = s * with_margin\n",
    "\n",
    "        loss = F.cross_entropy(with_margin, y)\n",
    "        return p, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArcTrainer(BaseEmbeddingTrainer):\n",
    "    def __init__(self, *, train_classes, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.accuracy = torchmetrics.Accuracy(\n",
    "            task=\"multiclass\",\n",
    "            num_classes=train_classes,\n",
    "        )\n",
    "\n",
    "        ### YOUR CODE HERE\n",
    "        self.classifier = ArcLoss(\n",
    "            EMB_DIM,\n",
    "            train_classes,\n",
    "        )\n",
    "\n",
    "    def training_step(self, batch):\n",
    "        x, y = batch\n",
    "\n",
    "        f = self.emb_net(x)\n",
    "        p, loss = self.classifier(f, y)\n",
    "        accs = self.accuracy(p.argmax(dim=1), y)\n",
    "\n",
    "        return self._log_metrics(loss, accs, \"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Зададим провайдеры данных."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = data.WeightedRandomSampler(\n",
    "    train_data.weights,\n",
    "    len(train_data.weights),\n",
    ")\n",
    "arc_train_dl = data.DataLoader(\n",
    "    train_data,\n",
    "    BATCH_SIZE,\n",
    "    sampler=sampler,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Обучим сеть."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_arc = train_model(\n",
    "    ArcTrainer(\n",
    "        emb_net=EmbeddingNet(),\n",
    "        train_classes=len(train_data.people_list),\n",
    "    ),\n",
    "    \"runs/arc\",\n",
    "    arc_train_dl,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Протестируем качество k-NN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_arc.eval()\n",
    "net_arc.emb_net.to(DEVICE)\n",
    "knn = train_knn(net_arc.emb_net, normalize=True)\n",
    "test_knn(net_arc.emb_net, knn, normalize=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}