{ "cells": [ { "cell_type": "markdown", "id": "475a2eaa", "metadata": {}, "source": [ "# Triplet 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_batches_per_epoch\"]` and `CONFIG_DATALOADER[\"val_batches\"]`: number of training and validation batches actually used per epoch.\n", "> - `CONFIG_TRAINER[\"max_epochs\"]`: number of epochs\n", "\n", "This notebook demonstrates how to train one `Triplet Neural Network` for image similarity computation. For this, the `Oxford Flowers` dataset is used, but customized for the triplet network training." ] }, { "cell_type": "markdown", "id": "bf8fb122", "metadata": {}, "source": [ "## Import libraries" ] }, { "cell_type": "code", "execution_count": null, "id": "f1f4f8a8", "metadata": {}, "outputs": [], "source": [ "from typing import Any\n", "from collections import defaultdict\n", "from collections.abc import Iterator\n", "import multiprocessing\n", "import random\n", "import datetime\n", "import os\n", "\n", "import torch\n", "from torch.utils.data import DataLoader, Dataset, Sampler\n", "from torch.utils.tensorboard import SummaryWriter\n", "from torchvision import transforms\n", "\n", "from pyvisim.neural_networks import TripletNeuralNetwork\n", "from pyvisim.datasets import OxfordFlowerDataset" ] }, { "cell_type": "markdown", "id": "c2637d26", "metadata": {}, "source": [ "## Declare the dataset\n", "\n", "A labeled batch of images goes through the shared-weight tower in a single pass, and the loss mines the (anchor, positive, negative) triplets from that batch itself.\n", "\n", "How a batch is drawn is very important. If drawn at random, a dataset with 1000 classes will yield very few positive pairs if the batch size is not large enough, and the network will learn very slowly. \n", "\n", "Hence, the `PKBatchSampler` is implemented. It draws `classes_per_batch` classes and `samples_per_class` images of each of them. `batch_size=classes_per_batch * samples_per_class`. With this, each batch is guaranteed to contain `samples_per_class - 1 ` positives and `batch_size - samples_per_class` negatives (assuming every sampled class has at least `samples_per_class` examples) for each image.\n", "\n", "The resulting batch is passed through the network, producing a `batch_size x batch_size` matrix of distances. The `TripletLoss` then mines the (hardest) triplets from this matrix and computes the loss. The network is trained to minimize this loss." ] }, { "cell_type": "code", "execution_count": null, "id": "bde73410", "metadata": {}, "outputs": [], "source": [ "class OxfordLabeledImages(Dataset[tuple[torch.Tensor, torch.Tensor]]):\n", " \"\"\"\n", " Labeled single images on top of :class:`OxfordFlowerDataset`.\n", "\n", " Each item is an ``(image, label)`` pair. The labels are also exposed as a\n", " plain list, so the batch sampler can group the indices by class without\n", " decoding a single image.\n", "\n", " :param purpose: Dataset split to draw from (e.g. ``\"train\"``).\n", " :param transform: Transform applied to each image.\n", " \"\"\"\n", "\n", " def __init__(self, purpose: str, transform: transforms.Compose) -> None:\n", " self._base = OxfordFlowerDataset(purpose=purpose)\n", " self.transform = transform\n", " self.labels: list[int] = list(self._base.labels)\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 __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:\n", " label = torch.tensor(self.labels[idx], dtype=torch.long)\n", " return self._get_tensor(idx), label" ] }, { "cell_type": "markdown", "id": "98176d77", "metadata": {}, "source": [ "### The P-K batch sampler\n", "\n", "`fixed_batches=True` restarts the sampler from the same seed on every epoch. The validation loss is then measured on the same triplets each time and stays comparable across epochs." ] }, { "cell_type": "code", "execution_count": null, "id": "635a8a31", "metadata": {}, "outputs": [], "source": [ "class PKBatchSampler(Sampler[list[int]]):\n", " \"\"\"\n", " P-K batch sampler: ``classes_per_batch`` classes, ``samples_per_class`` images each.\n", "\n", " Classes with fewer than ``samples_per_class`` images are sampled with\n", " replacement, so every batch has exactly the same size.\n", "\n", " :param labels: Class label of each dataset index.\n", " :param classes_per_batch: Number of distinct classes per batch (P).\n", " :param samples_per_class: Number of images drawn per sampled class (K).\n", " :param seed: Seed of the internal random generator.\n", " :param num_batches: Number of batches drawn per epoch.\n", " :param fixed_batches: If ``True``, every epoch yields the same batches.\n", " :raises ValueError: If the split holds fewer classes than\n", " ``classes_per_batch``, or if fewer than two classes are available.\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " labels: list[int],\n", " classes_per_batch: int,\n", " samples_per_class: int,\n", " seed: int,\n", " num_batches: int,\n", " fixed_batches: bool = False,\n", " ) -> None:\n", " # Labels are read directly, as indexing the base dataset would decode\n", " # every image just to group the indices by class.\n", " self._label_to_indices: dict[int, list[int]] = defaultdict(list)\n", " for idx, label in enumerate(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 triplet training.\")\n", " if len(self._labels_list) < classes_per_batch:\n", " raise ValueError(\n", " f\"Need >= {classes_per_batch} classes for a batch of \"\n", " f\"{classes_per_batch} classes, got {len(self._labels_list)}.\"\n", " )\n", "\n", " self.classes_per_batch = classes_per_batch\n", " self.samples_per_class = samples_per_class\n", " self.fixed_batches = fixed_batches\n", " self._seed = seed\n", " self._rng = random.Random(seed)\n", " self._num_batches = num_batches\n", "\n", " def __len__(self) -> int:\n", " return self._num_batches\n", "\n", " def _sample_class(self, label: int, rng: random.Random) -> list[int]:\n", " indices = self._label_to_indices[label]\n", " if len(indices) >= self.samples_per_class:\n", " return rng.sample(indices, self.samples_per_class)\n", " return rng.choices(indices, k=self.samples_per_class)\n", "\n", " def __iter__(self) -> Iterator[list[int]]:\n", " rng = random.Random(self._seed) if self.fixed_batches else self._rng\n", " for _ in range(self._num_batches):\n", " labels = rng.sample(self._labels_list, self.classes_per_batch)\n", " yield [idx for label in labels for idx in self._sample_class(label, rng)]" ] }, { "cell_type": "markdown", "id": "e3193a71", "metadata": {}, "source": [ "## Training seeds\n", "\n", "One pass over the training split takes `len(train_ds) // batch_size` batches." ] }, { "cell_type": "code", "execution_count": null, "id": "08e44b05", "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", "# The batch size is classes_per_batch * samples_per_class = 64. The smallest\n", "# split of Oxford Flowers still holds 10 images per class, so K = 4 fits into\n", "# every class without resampling.\n", "CONFIG_DATALOADER: dict[str, Any] = {\n", " \"classes_per_batch\": 16,\n", " \"samples_per_class\": 4,\n", " \"num_workers\": 4,\n", " \"train_batches_per_epoch\": 40,\n", " \"val_batches\": 12\n", "}\n", "\n", "def seed_worker(worker_id: int) -> None:\n", " random.seed(torch.initial_seed() % 2**32)" ] }, { "cell_type": "markdown", "id": "6f1b0b91", "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": "d791d31a", "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": "9816e92c", "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": "3646b477", "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 ``OxfordLabeledImages`` usable.\n", "MP_CONTEXT = \"fork\" if \"fork\" in multiprocessing.get_all_start_methods() else None\n", "\n", "def make_loader(\n", " dataset: OxfordLabeledImages,\n", " fixed_batches: bool,\n", " num_batches: int,\n", ") -> DataLoader[tuple[torch.Tensor, torch.Tensor]]:\n", " num_workers = CONFIG_DATALOADER[\"num_workers\"] if MP_CONTEXT is not None else 0\n", " # The sampler already shuffles, so the loader takes no \"shuffle\" flag.\n", " batch_sampler = PKBatchSampler(\n", " dataset.labels,\n", " classes_per_batch=CONFIG_DATALOADER[\"classes_per_batch\"],\n", " samples_per_class=CONFIG_DATALOADER[\"samples_per_class\"],\n", " seed=SEED,\n", " num_batches=num_batches,\n", " fixed_batches=fixed_batches,\n", " )\n", " return DataLoader(\n", " dataset,\n", " batch_sampler=batch_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": "e490c7ac", "metadata": {}, "source": [ "## Part 1: Semi-hard online mining\n", "\n", "Introduced by Hoffer and Ailon in 2014 and popularized by *FaceNet* (Schroff et al., 2015), the triplet network learns the embedding space from relative comparisons instead of absolute ones: an anchor has to be closer to a positive image (same class) than to a negative one (different class), by at least a margin.\n", "\n", "The three branches of the classical drawing are one single network. Anchor, positive and negative all pass through the same backbone (ResNet-18 here) and the same projection head, and the resulting embeddings are L2-normalized.\n", "\n", "`TripletLoss` is defined as follows:\n", "\n", "$$\n", "L(a, p, n) = \\max\\Big(0, \\, d(a, p) - d(a, n) + m\\Big)\n", "$$\n", "\n", "Whereas:\n", "\n", "- *a*, *p*, *n*: anchor, positive (same class as the anchor) and negative (different class)\n", "- *d*: the (optionally squared) Euclidean distance between two embeddings\n", "- *m*: margin: the minimum gap enforced between the positive and the negative distance\n", "\n", "The triplets are mined online, meaning that they are picked from the labeled batch during the forward pass. `mining=\"semi_hard\"` is FaceNet's strategy: for every positive pair, the closest negative that is still farther away than the positive is chosen. Such negatives already violate the margin and hence produce a gradient, without being so hard that the embedding collapses." ] }, { "cell_type": "code", "execution_count": null, "id": "a6ce776d", "metadata": {}, "outputs": [], "source": [ "from pyvisim.neural_networks.losses import TripletLoss" ] }, { "cell_type": "markdown", "id": "12242c38", "metadata": {}, "source": [ "### Train parameters" ] }, { "cell_type": "code", "execution_count": null, "id": "65871b52", "metadata": {}, "outputs": [], "source": [ "CONFIG_SEMI_HARD: dict[str, Any] = {\n", " \"embedding_dim\": 64,\n", " \"margin\": 0.2,\n", " \"mining\": \"semi_hard\",\n", " \"squared\": True,\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": "2a12908c", "metadata": {}, "source": [ "#### Some train helpers\n", "\n", "Both parts of this notebook train the very same network with the very same loop. Only the mining configuration changes, so the helpers below take the config and the run name as parameters." ] }, { "cell_type": "code", "execution_count": null, "id": "ae9fcae5", "metadata": {}, "outputs": [], "source": [ "def _run_epoch_triplet(\n", " model: TripletNeuralNetwork,\n", " loader: DataLoader[tuple[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", " config: dict[str, Any],\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, (images, labels) in enumerate(loader):\n", " images = images.to(device)\n", " labels = labels.to(device)\n", "\n", " # One shared-weight pass over the whole batch. The triplets are\n", " # mined from these embeddings by the loss itself.\n", " embeddings = model(images)\n", " loss = criterion(embeddings, 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[\"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_triplet(\n", " model: TripletNeuralNetwork,\n", " criterion: torch.nn.Module,\n", " config: dict[str, Any],\n", " run_name: str,\n", ") -> None:\n", "\n", " session_id = datetime.datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n", " writer = SummaryWriter(log_dir=f\"runs/triplet/session_{session_id}_{run_name}\")\n", "\n", " try:\n", " os.makedirs(config[\"checkpoint_dir\"], exist_ok=True)\n", "\n", " train_ds = OxfordLabeledImages(\"train\", transform=TRAIN_TF)\n", " val_ds = OxfordLabeledImages(\"validation\", transform=VAL_TF)\n", "\n", " train_loader = make_loader(\n", " train_ds, fixed_batches=False, num_batches=CONFIG_DATALOADER[\"train_batches_per_epoch\"]\n", " )\n", " val_loader = make_loader(\n", " val_ds, fixed_batches=True, num_batches=CONFIG_DATALOADER[\"val_batches\"]\n", " )\n", "\n", " batch_size = (\n", " CONFIG_DATALOADER[\"classes_per_batch\"]\n", " * CONFIG_DATALOADER[\"samples_per_class\"]\n", " )\n", " print(f\"Train images: {len(train_ds):,} | Val images: {len(val_ds):,}\")\n", " print(f\"Flower classes: {len(set(train_ds.labels))}\")\n", " print(f\"Batch size: {batch_size} | Batches per epoch: {len(train_loader)}\")\n", "\n", " optimizer = torch.optim.AdamW(\n", " model.parameters(), lr=config[\"lr\"], weight_decay=config[\"weight_decay\"]\n", " )\n", " scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=config.get('gamma', 0.8))\n", "\n", " best_val = float(\"inf\")\n", " for epoch in range(1, config[\"num_epochs\"] + 1):\n", " train_loss = _run_epoch_triplet(\n", " model, train_loader, criterion, optimizer, is_train=True, epoch=epoch, writer=writer, config=config\n", " )\n", " val_loss = _run_epoch_triplet(\n", " model, val_loader, criterion, optimizer, is_train=False, epoch=epoch, writer=writer, config=config\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(f\"{run_name}.safetensors\")\n", " if val_loss < best_val:\n", " best_val = val_loss\n", " model.save_to_disk(f\"{run_name}_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": "8aa80477", "metadata": {}, "outputs": [], "source": [ "criterion = TripletLoss(\n", " margin=CONFIG_SEMI_HARD[\"margin\"],\n", " mining=CONFIG_SEMI_HARD[\"mining\"],\n", " squared=CONFIG_SEMI_HARD[\"squared\"],\n", ")\n", "model_semi_hard = TripletNeuralNetwork(\n", " backbone=\"resnet18\",\n", " embedding_dim=CONFIG_SEMI_HARD[\"embedding_dim\"],\n", " device=device,\n", " pretrained_backbone=True,\n", ")\n", "train_triplet(model_semi_hard, criterion, CONFIG_SEMI_HARD, run_name=\"triplet_semi_hard\")" ] }, { "cell_type": "markdown", "id": "45996132", "metadata": {}, "source": [ "### Load models after training" ] }, { "cell_type": "code", "execution_count": null, "id": "fcfcb5c2", "metadata": {}, "outputs": [], "source": [ "model_semi_hard = TripletNeuralNetwork.load_from_disk(\"triplet_semi_hard_best.safetensors\")\n", "model_semi_hard.eval()" ] }, { "cell_type": "markdown", "id": "31668496", "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": "8c92cd46", "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": "1c758ee0", "metadata": {}, "source": [ "### Compute similarity\n", "\n", "As observed, the pair of the same class yields a similarity score close to 1, while the pair of different classes yields a similarity score should be much lower (at least below 0.0, or at best close to -1.0). The embeddings are L2-normalized, so the cosine similarity is simply their dot product." ] }, { "cell_type": "code", "execution_count": null, "id": "5b29729a", "metadata": {}, "outputs": [], "source": [ "score_same_class = model_semi_hard.similarity_score(img_a, img_b).item()\n", "score_different_class = model_semi_hard.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": "cfa7dcef", "metadata": {}, "source": [ "## Part 2: Batch-hard online mining\n", "\n", "*In Defense of the Triplet Loss* (Hermans et al., 2017)'s approach keeps only the two extremes of every anchor inside the batch: its farthest positive and its closest negative. Fewer triplets contribute, but each of them carries a much stronger gradient. The same paper reports using Euclidean distance is more stable, while the squared distance made the optimization more prone to collapsing (all embeddings mapped to the same point), which is why `squared` is turned off here.\n", "\n", "### NOTE\n", "\n", "This is the very same network and the very same training loop as in part 1. Only the mined triplets differ, which is also why the P-K sampling matters even more here: an anchor without a positive in its batch contributes nothing at all to the batch-hard loss." ] }, { "cell_type": "markdown", "id": "1523680c", "metadata": {}, "source": [ "### Train parameters" ] }, { "cell_type": "code", "execution_count": null, "id": "e3eec075", "metadata": {}, "outputs": [], "source": [ "CONFIG_BATCH_HARD: dict[str, Any] = {\n", " \"embedding_dim\": 64,\n", " \"margin\": 0.5,\n", " \"mining\": \"batch_hard\",\n", " \"squared\": False,\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": "code", "execution_count": null, "id": "75c25c1a", "metadata": {}, "outputs": [], "source": [ "criterion = TripletLoss(\n", " margin=CONFIG_BATCH_HARD[\"margin\"],\n", " mining=CONFIG_BATCH_HARD[\"mining\"],\n", " squared=CONFIG_BATCH_HARD[\"squared\"],\n", ")\n", "model_batch_hard = TripletNeuralNetwork(\n", " backbone=\"resnet18\",\n", " embedding_dim=CONFIG_BATCH_HARD[\"embedding_dim\"],\n", " device=device,\n", " pretrained_backbone=True,\n", ")\n", "train_triplet(model_batch_hard, criterion, CONFIG_BATCH_HARD, run_name=\"triplet_batch_hard\")" ] }, { "cell_type": "markdown", "id": "74a05b60", "metadata": {}, "source": [ "### Load models after training" ] }, { "cell_type": "code", "execution_count": null, "id": "0e2963ea", "metadata": {}, "outputs": [], "source": [ "model_batch_hard = TripletNeuralNetwork.load_from_disk(\"triplet_batch_hard_best.safetensors\")\n", "model_batch_hard.eval()" ] }, { "cell_type": "markdown", "id": "bd80c3f4", "metadata": {}, "source": [ "### Visualize results\n", "\n", "The same three images are reused, so the two mining strategies can be compared on identical inputs." ] }, { "cell_type": "code", "execution_count": null, "id": "1e4f749e", "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": "b93405a9", "metadata": {}, "source": [ "### Compute similarity\n", "\n", "Both networks produce L2-normalized embeddings, so the scores below live on the same scale as the ones of part 1 and can be read side by side." ] }, { "cell_type": "code", "execution_count": null, "id": "2eb98e43", "metadata": {}, "outputs": [], "source": [ "score_same_class = model_batch_hard.similarity_score(img_a, img_b).item()\n", "score_different_class = model_batch_hard.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": "785904dd", "metadata": {}, "source": [ "## References\n", "\n", "[1] Hoffer, E., & Ailon, N. (2014). Deep Metric Learning Using Triplet\n", "Network. https://arxiv.org/abs/1412.6622\n", "\n", "\n", "[2] Schroff, F., Kalenichenko, D., & Philbin, J. (2015). FaceNet: A Unified\n", "Embedding for Face Recognition and Clustering. In Proceedings of the 2015\n", "IEEE Conference on Computer Vision and Pattern Recognition (CVPR),\n", "815-823. https://doi.org/10.1109/CVPR.2015.7298682\n", "\n", "\n", "[3] Hermans, A., Beyer, L., & Leibe, B. (2017). In Defense of the Triplet\n", "Loss for Person Re-Identification. https://arxiv.org/abs/1703.07737" ] } ], "metadata": {}, "nbformat": 4, "nbformat_minor": 5 }