{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba749009",
   "metadata": {},
   "source": [
    "# Siamese Neural Network\n",
    "\n",
    "> [!NOTE]\n",
    "> The number of epochs and the real dataset size actually used for training and validation is reduced to speed up the CI pipeline. To improve th performance of the model, feel free to increase these:\n",
    "> - `CONFIG_DATALOADER[\"train_pairs_per_epoch\"]` and `CONFIG_DATALOADER[\"val_pairs\"]`: number of training and validation pairs actually used per epoch.\n",
    "> - `CONFIG_TRAINER[\"max_epochs\"]`: number of epochs\n",
    "\n",
    "This notebook demonstrates how to train one `Siamese Neural Network` for image similarity computation. For this,  the `Oxford Flowers` dataset is used, but customized for the Siamese network training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4ce0355",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a310c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "from collections import defaultdict\n",
    "import multiprocessing\n",
    "import random\n",
    "import datetime\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchvision import transforms\n",
    "\n",
    "from pyvisim.neural_networks import ContrastiveSiameseNetwork, BCESiameseNetwork\n",
    "from pyvisim.datasets import OxfordFlowerDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e84808",
   "metadata": {},
   "source": [
    "## Declare the dataset\n",
    "\n",
    "The `Siamese Network` needs pairs of images during training. Each item in the dataset is a pair of images (img_a, img_b) and a label indicating whether they are `similar` (interger 1) or `dissimilar` (integer 0).\n",
    "\n",
    "The pairs are also built upfront to avoid repeated sampling during training. You can change the seed to get different pairs.\n",
    "\n",
    "`positive_fraction` controls the ratio of positive to negative pairs generated during training. For example, if `positive_fraction=0.5`, then half of the pairs will be positive and half will be negative. The default ratio of `0.5` assumes a balanced dataset. For skewed datasets, adjust this parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "870347ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class OxfordSiamesePairs(Dataset[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]):\n",
    "    def __init__(\n",
    "        self,\n",
    "        purpose: str,\n",
    "        transform: transforms.Compose,\n",
    "        positive_fraction: float = 0.5,\n",
    "    ) -> None:\n",
    "        self._base = OxfordFlowerDataset(purpose=purpose)\n",
    "        self.transform = transform\n",
    "        self.positive_fraction = positive_fraction\n",
    "\n",
    "        # Build a label -> list[index] map for fast pair mining. Labels are\n",
    "        # read directly; indexing the base dataset would decode every image.\n",
    "        self._label_to_indices: dict[int, list[int]] = defaultdict(list)\n",
    "        for idx, label in enumerate(self._base.labels):\n",
    "            self._label_to_indices[label].append(idx)\n",
    "\n",
    "        self._labels_list = list(self._label_to_indices.keys())\n",
    "        if len(self._labels_list) < 2:\n",
    "            raise ValueError(\"Need >= 2 classes for contrastive training.\")\n",
    "\n",
    "        # For reproducibility\n",
    "        rng = random.Random(SEED)\n",
    "\n",
    "        # Build the pairs list upfront to avoid repeated sampling during training, \n",
    "        # which has O(n) complexity.\n",
    "        self._pairs: list[tuple[int, int]] = [\n",
    "            self._sample_partner(idx, rng) for idx in range(len(self._base))\n",
    "        ]\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self._base)\n",
    "\n",
    "    def _get_tensor(self, idx: int) -> torch.Tensor:\n",
    "        image, _, _ = self._base[idx]\n",
    "        tensor: torch.Tensor = self.transform(image)\n",
    "        return tensor\n",
    "\n",
    "    def _sample_partner(self, idx: int, rng) -> tuple[int, int]:\n",
    "        label_a = self._base.labels[idx]\n",
    "        if rng.random() < self.positive_fraction:\n",
    "            # Positive: another image of the same class.\n",
    "            same_indices = [i for i in self._label_to_indices[label_a] if i != idx]\n",
    "            idx_b = rng.choice(same_indices) if same_indices else idx\n",
    "            return idx_b, 1\n",
    "        # Negative: an image from a different class.\n",
    "        neg_label = rng.choice([lbl for lbl in self._labels_list if lbl != label_a])\n",
    "        idx_b = rng.choice(self._label_to_indices[neg_label])\n",
    "        return idx_b, 0\n",
    "\n",
    "    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        img_a = self._get_tensor(idx)\n",
    "        idx_b, pair_label = self._pairs[idx]\n",
    "        img_b = self._get_tensor(idx_b)\n",
    "        return img_a, img_b, torch.tensor(pair_label, dtype=torch.float32)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64af8ddf",
   "metadata": {},
   "source": [
    "## Training seeds\n",
    "\n",
    "Training on every pair of the training split for many epochs takes hours without a GPU. `train_pairs_per_epoch` draws that many random pairs in every epoch and `val_pairs` validates on the first pairs of the validation split, so that this notebook also finishes on a CPU. The splits hold one pair per image, so raise both up to the size of the splits to train on every pair."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2fb9efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "CONFIG_DATALOADER: dict[str, Any] = {\n",
    "    \"batch_size\": 64,\n",
    "    \"num_workers\": 4,\n",
    "    \"train_pairs_per_epoch\": 640,\n",
    "    \"val_pairs\": 160\n",
    "}\n",
    "\n",
    "def seed_worker(worker_id: int) -> None:\n",
    "    random.seed(torch.initial_seed() % 2**32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5644989a",
   "metadata": {},
   "source": [
    "## Transforms\n",
    "\n",
    "The transforms below match the ones used in the pre-trained backbone (*ResNet-18*, in the default configuration, which is also the configuration below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876e5a0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "_NORMALIZE = transforms.Normalize(\n",
    "    mean=[0.485, 0.456, 0.406],\n",
    "    std=[0.229, 0.224, 0.225],\n",
    ")\n",
    "\n",
    "TRAIN_TF = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.RandomApply(\n",
    "            [\n",
    "                transforms.RandomChoice([\n",
    "                    transforms.RandomHorizontalFlip(p=1.0),\n",
    "                    transforms.RandomRotation(15),\n",
    "                ]),\n",
    "            ],\n",
    "            p=0.5,\n",
    "        ),\n",
    "        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1),\n",
    "        transforms.ToTensor(),\n",
    "        _NORMALIZE,\n",
    "    ]\n",
    ")\n",
    "\n",
    "VAL_TF = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        _NORMALIZE,\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba0f7f50",
   "metadata": {},
   "source": [
    "## Multiprocessing configs\n",
    "\n",
    "This was necessary on `Linux` and `Python 3.14+`. You can comment it out if it caused you trouble."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bb779cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python 3.14 made \"forkserver\" the default start method on Linux. A forkserver\n",
    "# worker re-imports the payload in a fresh interpreter, so it cannot resolve\n",
    "# classes defined in a notebook (they live in ``__main__``). Forking inherits the\n",
    "# interpreter state instead, which keeps ``OxfordSiamesePairs`` usable.\n",
    "MP_CONTEXT = \"fork\" if \"fork\" in multiprocessing.get_all_start_methods() else None\n",
    "\n",
    "def make_loader(\n",
    "    dataset: Dataset[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],\n",
    "    num_samples: int,\n",
    "    shuffle: bool,\n",
    ") -> DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:\n",
    "    num_workers = CONFIG_DATALOADER[\"num_workers\"] if MP_CONTEXT is not None else 0\n",
    "    # A shuffled loader draws new random items in every epoch, the other one\n",
    "    # always reads the first items of the dataset.\n",
    "    sampler = (\n",
    "        RandomSampler(dataset, num_samples=num_samples)\n",
    "        if shuffle\n",
    "        else SequentialSampler(range(min(num_samples, len(dataset))))\n",
    "    )\n",
    "    return DataLoader(\n",
    "        dataset,\n",
    "        batch_size=CONFIG_DATALOADER[\"batch_size\"],\n",
    "        sampler=sampler,\n",
    "        num_workers=num_workers,\n",
    "        pin_memory=True,\n",
    "        worker_init_fn=seed_worker,\n",
    "        multiprocessing_context=MP_CONTEXT if num_workers > 0 else None,\n",
    "        persistent_workers=num_workers > 0,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12388de2",
   "metadata": {},
   "source": [
    "## Part 1: The contrastive Siamese network\n",
    "\n",
    "Introduced by LeCun et al. in 2005, this was the classical metric-learning approach for image similarity. The network is trained to minimize the distance between embeddings of similar images and maximize the distance between embeddings of dissimilar images.\n",
    "\n",
    "First, it encodes the input images into embeddings using a backbone network (ResNet-18 in this case). Then, it computes the distance between the embeddings. The `ContrastiveLoss` is used as criterion.\n",
    "\n",
    "`ContrastiveLoss` is defined as follows:\n",
    "\n",
    "$$\n",
    "L = \\frac{1}{2N} \\sum_{n=1}^{N} \\Big[ y_n d_n^2 + (1 - y_n) \\max(0, m - d_n)^2 \\Big]\n",
    "$$\n",
    "\n",
    "Whereas: \n",
    "\n",
    "- *N*: number of pairs in the batch\n",
    "- *y_n*: label of the pair (1 for similar, 0 for dissimilar)\n",
    "- *d_n*: distance between the embeddings of the pair\n",
    "- *m*: margin: controls how far apart the embeddings of dissimilar pairs should be"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a3b3306",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyvisim.neural_networks.losses import ContrastiveLoss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4873ca8",
   "metadata": {},
   "source": [
    "### Train parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0fd769f",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_CONTRASTIVE: dict[str, Any] = {\n",
    "    \"embedding_dim\": 64,\n",
    "    \"margin\": 1.0,\n",
    "    \"lr\": 1e-4,\n",
    "    \"weight_decay\": 1e-4,\n",
    "    \"num_epochs\": 5,\n",
    "    \"checkpoint_dir\": \"checkpoints\",\n",
    "    \"log_every_n\": 50,\n",
    "    \"gamma\": 0.8, \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080005e2",
   "metadata": {},
   "source": [
    "#### Some train helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2f52999",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _run_epoch_contrastive(\n",
    "    model: ContrastiveSiameseNetwork,\n",
    "    loader: DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],\n",
    "    criterion: torch.nn.Module,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    is_train: bool,\n",
    "    epoch: int,\n",
    "    writer: SummaryWriter\n",
    ") -> float:\n",
    "    model.train() if is_train else model.eval()\n",
    "    total_loss = 0.0\n",
    "    ctx = torch.enable_grad() if is_train else torch.no_grad()\n",
    "\n",
    "    with ctx:\n",
    "        for step, (img_a, img_b, labels) in enumerate(loader):\n",
    "            img_a = img_a.to(device)\n",
    "            img_b = img_b.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            emb_a = model(img_a)\n",
    "            emb_b = model(img_b)\n",
    "            loss = criterion(emb_a, emb_b, labels)\n",
    "\n",
    "            if is_train:\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                # Gradient clipping to prevent exploding gradients.\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                optimizer.step()\n",
    "\n",
    "            total_loss += loss.item()\n",
    "\n",
    "            if is_train and (step + 1) % CONFIG_CONTRASTIVE[\"log_every_n\"] == 0:\n",
    "                print(\n",
    "                    f\"  Epoch {epoch:03d} | step {step + 1:04d} \"\n",
    "                    f\"| loss {total_loss / (step + 1):.4f}\"\n",
    "                )\n",
    "\n",
    "            loss_str = \"Loss/Train\" if is_train else \"Loss/Validation\"\n",
    "            global_step = (epoch - 1) * len(loader) + step\n",
    "            writer.add_scalar(loss_str, loss.item(), global_step=global_step)\n",
    "\n",
    "    return total_loss / len(loader)\n",
    "\n",
    "\n",
    "def train_contrastive(model: ContrastiveSiameseNetwork, criterion: torch.nn.Module) -> None:\n",
    "\n",
    "    session_id = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    writer = SummaryWriter(log_dir=f\"runs/siamese/session_{session_id}_contrastiveloss\")\n",
    "\n",
    "    try:\n",
    "        os.makedirs(CONFIG_CONTRASTIVE[\"checkpoint_dir\"], exist_ok=True)\n",
    "\n",
    "        train_ds = OxfordSiamesePairs(\"train\", transform=TRAIN_TF)\n",
    "        val_ds = OxfordSiamesePairs(\"validation\", transform=VAL_TF)\n",
    "\n",
    "        train_loader = make_loader(train_ds, CONFIG_DATALOADER[\"train_pairs_per_epoch\"], shuffle=True)\n",
    "        val_loader = make_loader(val_ds, CONFIG_DATALOADER[\"val_pairs\"], shuffle=False)\n",
    "\n",
    "        print(\n",
    "            f\"Train pairs per epoch: {len(train_loader.sampler):,} of {len(train_ds):,}  \"\n",
    "            f\"|  Val pairs: {len(val_loader.sampler):,}\"\n",
    "        )\n",
    "        print(f\"Flower classes: {len(train_ds._labels_list)}\")\n",
    "\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            model.parameters(), lr=CONFIG_CONTRASTIVE[\"lr\"], weight_decay=CONFIG_CONTRASTIVE[\"weight_decay\"]\n",
    "        )\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=CONFIG_CONTRASTIVE.get('gamma', 0.8))\n",
    "\n",
    "        best_val = float(\"inf\")\n",
    "        for epoch in range(1, CONFIG_CONTRASTIVE[\"num_epochs\"] + 1):\n",
    "            train_loss = _run_epoch_contrastive(\n",
    "                model, train_loader, criterion, optimizer, is_train=True, epoch=epoch, writer=writer\n",
    "            )\n",
    "            val_loss = _run_epoch_contrastive(\n",
    "                model, val_loader, criterion, optimizer, is_train=False, epoch=epoch, writer=writer\n",
    "            )\n",
    "            scheduler.step()\n",
    "\n",
    "            print(\n",
    "                f\"Epoch {epoch:03d} | train {train_loss:.4f} | val {val_loss:.4f} \"\n",
    "                f\"| lr {scheduler.get_last_lr()[0]:.2e}\"\n",
    "            )\n",
    "            model.save_to_disk(\"contrastive_siamese.safetensors\")\n",
    "            if val_loss < best_val:\n",
    "                best_val = val_loss\n",
    "                model.save_to_disk(\"contrastive_siamese_best.safetensors\")\n",
    "                print(f\"New best val loss: {best_val:.4f}\")\n",
    "\n",
    "        print(f\"\\nTraining complete. Best val loss: {best_val:.4f}\")\n",
    "    finally:\n",
    "        writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37c489e",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = ContrastiveLoss(margin=CONFIG_CONTRASTIVE[\"margin\"])\n",
    "model_contrastive = ContrastiveSiameseNetwork(\n",
    "    backbone=\"resnet18\",\n",
    "    embedding_dim=CONFIG_CONTRASTIVE[\"embedding_dim\"],\n",
    "    device=device,\n",
    "    pretrained_backbone=True,\n",
    ")\n",
    "train_contrastive(model_contrastive, criterion)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ab54723",
   "metadata": {},
   "source": [
    "### Load models after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f7f8bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_contrastive = ContrastiveSiameseNetwork.load_from_disk(\"contrastive_siamese_best.safetensors\")\n",
    "model_contrastive.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca4ff0ea",
   "metadata": {},
   "source": [
    "### Visualize results\n",
    "\n",
    "Now, we will pick three random images from the validation set, two of which are similar and one is dissimilar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76a07b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def visualize_image_pair(image_a: np.ndarray, image_b: np.ndarray) -> None:\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
    "    axes[0].imshow(image_a)\n",
    "    axes[0].set_title(\"Image A\")\n",
    "    axes[0].axis(\"off\")\n",
    "\n",
    "    axes[1].imshow(image_b)\n",
    "    axes[1].set_title(\"Image B\")\n",
    "    axes[1].axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "test_dataset =  OxfordFlowerDataset(purpose=\"test\")\n",
    "img_a = test_dataset[4][0]\n",
    "img_b = test_dataset[5][0]\n",
    "img_c = test_dataset[100][0]\n",
    "print(\"Image shapes:\\n\")\n",
    "print(f\"Image A: {img_a.shape}\")\n",
    "print(f\"Image B: {img_b.shape}\")\n",
    "\n",
    "visualize_image_pair(img_a, img_b)\n",
    "visualize_image_pair(img_a, img_c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4992832b",
   "metadata": {},
   "source": [
    "### Compute similarity \n",
    "\n",
    "As observed, the pair of the same class yields a similarity score closer to 1, while the pair of different classes yields a similarity score should be much lower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "281c4ab8",
   "metadata": {},
   "outputs": [],
   "source": [
    "score_same_class = model_contrastive.similarity_score(img_a, img_b).item()\n",
    "score_different_class = model_contrastive.similarity_score(img_a, img_c).item()\n",
    "\n",
    "print(f\"Similarity score (same class): {score_same_class:.4f}\")\n",
    "print(f\"Similarity score (different class): {score_different_class:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f701f819",
   "metadata": {},
   "source": [
    "## Part 2: BCELoss Siamese network\n",
    "\n",
    "This is a more recent approach to train a Siamese network. Instead of learning the metric, the network learns to predict the similarity score directly using `L1` distance between the embeddings and a `sigmoid` activation function. The network is trained using `BCELoss`, as familiar in binary classification tasks.\n",
    "\n",
    "### NOTE\n",
    "\n",
    "This network, unlike the contrastive network, does not learn to produce embeddings. Instead, they learn to model the similarity score directly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e66dad83",
   "metadata": {},
   "source": [
    "### Import the loss function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "107c451f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import BCEWithLogitsLoss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a684ba9",
   "metadata": {},
   "source": [
    "### Train parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23d5cbc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_BCE: dict[str, Any] = {\n",
    "    \"embedding_dim\": 128,\n",
    "    \"margin\": 1.0,\n",
    "    \"lr\": 1e-4,\n",
    "    \"weight_decay\": 1e-4,\n",
    "    \"num_epochs\": 5,\n",
    "    \"checkpoint_dir\": \"checkpoints\",\n",
    "    \"log_every_n\": 50,\n",
    "    \"gamma\": 0.8, \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5577424f",
   "metadata": {},
   "source": [
    "#### Train helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "171b6113",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _run_epoch_bce(\n",
    "    model: BCESiameseNetwork,\n",
    "    loader: DataLoader[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],\n",
    "    criterion: torch.nn.BCEWithLogitsLoss,\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    is_train: bool,\n",
    "    epoch: int,\n",
    ") -> float:\n",
    "    model.train() if is_train else model.eval()\n",
    "    total_loss = 0.0\n",
    "    ctx = torch.enable_grad() if is_train else torch.no_grad()\n",
    "\n",
    "    with ctx:\n",
    "        for step, (img_a, img_b, labels) in enumerate(loader):\n",
    "            img_a = img_a.to(device)\n",
    "            img_b = img_b.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            logits = model(img_a, img_b)\n",
    "            loss = criterion(logits, labels)\n",
    "\n",
    "            if is_train:\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                # Gradient clipping to prevent exploding gradients.\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n",
    "                optimizer.step()\n",
    "\n",
    "            total_loss += loss.item()\n",
    "\n",
    "            if is_train and (step + 1) % CONFIG_BCE[\"log_every_n\"] == 0:\n",
    "                print(\n",
    "                    f\"  Epoch {epoch:03d} | step {step + 1:04d} \"\n",
    "                    f\"| loss {total_loss / (step + 1):.4f}\"\n",
    "                )\n",
    "\n",
    "    return total_loss / len(loader)\n",
    "\n",
    "\n",
    "def train_bce(model: BCESiameseNetwork, criterion: torch.nn.BCEWithLogitsLoss) -> None:\n",
    "    session_id = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    writer = SummaryWriter(log_dir=f\"runs/siamese/session_{session_id}_bceloss\")\n",
    "\n",
    "    try:\n",
    "        os.makedirs(CONFIG_BCE[\"checkpoint_dir\"], exist_ok=True)\n",
    "\n",
    "        train_ds = OxfordSiamesePairs(\"train\", transform=TRAIN_TF)\n",
    "        val_ds = OxfordSiamesePairs(\"validation\", transform=VAL_TF)\n",
    "\n",
    "        train_loader = make_loader(train_ds, CONFIG_DATALOADER[\"train_pairs_per_epoch\"], shuffle=True)\n",
    "        val_loader = make_loader(val_ds, CONFIG_DATALOADER[\"val_pairs\"], shuffle=False)\n",
    "\n",
    "        print(\n",
    "            f\"Train pairs per epoch: {len(train_loader.sampler):,} of {len(train_ds):,}  \"\n",
    "            f\"|  Val pairs: {len(val_loader.sampler):,}\"\n",
    "        )\n",
    "        print(f\"Flower classes: {len(train_ds._labels_list)}\")\n",
    "\n",
    "        model = BCESiameseNetwork(\n",
    "            backbone=\"resnet18\",\n",
    "            embedding_dim=CONFIG_BCE[\"embedding_dim\"],\n",
    "            device=device,\n",
    "            pretrained_backbone=True,\n",
    "        )\n",
    "\n",
    "        criterion = torch.nn.BCEWithLogitsLoss()\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            model.parameters(), lr=CONFIG_BCE[\"lr\"], weight_decay=CONFIG_BCE[\"weight_decay\"]\n",
    "        )\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(\n",
    "            optimizer, \n",
    "            step_size=4, \n",
    "            gamma=CONFIG_BCE.get('gamma', 0.8)\n",
    "        )\n",
    "\n",
    "        best_val = float(\"inf\")\n",
    "        for epoch in range(1, CONFIG_BCE[\"num_epochs\"] + 1):\n",
    "            train_loss = _run_epoch_bce(\n",
    "                model, train_loader, criterion, optimizer, is_train=True, epoch=epoch\n",
    "            )\n",
    "            val_loss = _run_epoch_bce(\n",
    "                model, val_loader, criterion, optimizer, is_train=False, epoch=epoch\n",
    "            )\n",
    "            scheduler.step()\n",
    "\n",
    "            print(\n",
    "                f\"Epoch {epoch:03d} | train {train_loss:.4f} | val {val_loss:.4f} \"\n",
    "                f\"| lr {scheduler.get_last_lr()[0]:.2e}\"\n",
    "            )\n",
    "\n",
    "            writer.add_scalar(\"Loss/Train\", train_loss, epoch)\n",
    "            writer.add_scalar(\"Loss/Validation\", val_loss, epoch)\n",
    "            writer.add_scalar(\"Learning Rate\", scheduler.get_last_lr()[0], epoch)\n",
    "\n",
    "            ckpt = {\n",
    "                \"epoch\": epoch,\n",
    "                \"model\": model.state_dict(),\n",
    "                \"optimizer\": optimizer.state_dict(),\n",
    "                \"scheduler\": scheduler.state_dict(),\n",
    "                \"val_loss\": val_loss,\n",
    "                \"cfg\": CONFIG_BCE,\n",
    "            }\n",
    "            model.save_to_disk(\"bce_siamese.safetensors\")\n",
    "            if val_loss < best_val:\n",
    "                best_val = val_loss\n",
    "                model.save_to_disk(\"bce_siamese_best.safetensors\")\n",
    "                print(f\"New best val loss: {best_val:.4f}\")\n",
    "\n",
    "        print(f\"\\nTraining complete. Best val loss: {best_val:.4f}\")\n",
    "    finally:\n",
    "        writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "510907f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = BCEWithLogitsLoss()\n",
    "model_bce = BCESiameseNetwork(\n",
    "    backbone=\"resnet18\",\n",
    "    embedding_dim=CONFIG_BCE[\"embedding_dim\"],\n",
    "    device=device,\n",
    "    pretrained_backbone=True,\n",
    ")\n",
    "train_bce(model_bce, criterion)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2720d49",
   "metadata": {},
   "source": [
    "### Load models after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625f7c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_bce = BCESiameseNetwork.load_from_disk(\"bce_siamese_best.safetensors\")\n",
    "model_bce.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f875a153",
   "metadata": {},
   "source": [
    "### Visualize results\n",
    "\n",
    "Now, we will pick three random images from the validation set, two of which are similar and one is dissimilar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7b107d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Image shapes:\\n\")\n",
    "print(f\"Image A: {img_a.shape}\")\n",
    "print(f\"Image B: {img_b.shape}\")\n",
    "\n",
    "visualize_image_pair(img_a, img_b)\n",
    "visualize_image_pair(img_a, img_c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5f6bef2",
   "metadata": {},
   "source": [
    "### Compute similarity\n",
    "\n",
    " Since a sigmoid function is applied on the absolute difference between embeddings, the score for the similar pair should be closer to 1 and the score for the dissimilar pair should be closer to 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91974956",
   "metadata": {},
   "outputs": [],
   "source": [
    "score_same_class = model_bce.similarity_score(img_a, img_b).item()\n",
    "score_different_class = model_bce.similarity_score(img_a, img_c).item()\n",
    "\n",
    "print(f\"Similarity score (same class): {score_same_class:.4f}\")\n",
    "print(f\"Similarity score (different class): {score_different_class:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c1c9149",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "[1] Koch, G., Zemel, R., & Salakhutdinov, R. (2015). Siamese Neural Networks\n",
    "for One-shot Image Recognition. ICML Deep Learning Workshop.\n",
    "https://www.cs.cmu.edu/~rsalakhu/papers/oneshot1.pdf\n",
    "\n",
    "\n",
    "[2] Hadsell, R., Chopra, S., & LeCun, Y. (2006). Dimensionality Reduction\n",
    "by Learning an Invariant Mapping. In Proceedings of the 2006 IEEE\n",
    "Computer Society Conference on Computer Vision and Pattern Recognition\n",
    "(CVPR), Vol. 2, 1735-1742. https://doi.org/10.1109/CVPR.2006.100\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}