Metric learning for image similarity search

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/vision/metric_learning/" />  <meta property="og:title" content="Keras documentation: Metric learning for image similarity search"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Metric learning for image similarity search"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Metric learning for image similarity search</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink active" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink2" href="/examples/vision/image_classification_from_scratch/">Image classification from scratch</a> <a class="nav-sublink2" href="/examples/vision/mnist_convnet/">Simple MNIST convnet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_efficientnet_fine_tuning/">Image classification via fine-tuning with EfficientNet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_with_vision_transformer/">Image classification with Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/attention_mil_classification/">Classification using Attention-based Deep Multiple Instance Learning</a> <a class="nav-sublink2" href="/examples/vision/mlp_image_classification/">Image classification with modern MLP models</a> <a class="nav-sublink2" href="/examples/vision/mobilevit/">A mobile-friendly Transformer-based model for image classification</a> <a class="nav-sublink2" href="/examples/vision/xray_classification_with_tpus/">Pneumonia Classification on TPU</a> <a class="nav-sublink2" href="/examples/vision/cct/">Compact Convolutional Transformers</a> <a class="nav-sublink2" href="/examples/vision/convmixer/">Image classification with ConvMixer</a> <a class="nav-sublink2" href="/examples/vision/eanet/">Image classification with EANet (External Attention Transformer)</a> <a class="nav-sublink2" href="/examples/vision/involution/">Involutional neural networks</a> <a class="nav-sublink2" href="/examples/vision/perceiver_image_classification/">Image classification with Perceiver</a> <a class="nav-sublink2" href="/examples/vision/reptile/">Few-Shot learning with Reptile</a> <a class="nav-sublink2" href="/examples/vision/semisupervised_simclr/">Semi-supervised image classification using contrastive pretraining with SimCLR</a> <a class="nav-sublink2" href="/examples/vision/swin_transformers/">Image classification with Swin Transformers</a> <a class="nav-sublink2" href="/examples/vision/vit_small_ds/">Train a Vision Transformer on small datasets</a> <a class="nav-sublink2" href="/examples/vision/shiftvit/">A Vision Transformer without Attention</a> <a class="nav-sublink2" href="/examples/vision/image_classification_using_global_context_vision_transformer/">Image Classification using Global Context Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/temporal_latent_bottleneck/">When Recurrence meets Transformers</a> <a class="nav-sublink2" href="/examples/vision/oxford_pets_image_segmentation/">Image segmentation with a U-Net-like architecture</a> <a class="nav-sublink2" href="/examples/vision/deeplabv3_plus/">Multiclass semantic segmentation using DeepLabV3+</a> <a class="nav-sublink2" href="/examples/vision/basnet_segmentation/">Highly accurate boundaries segmentation using BASNet</a> <a class="nav-sublink2" href="/examples/vision/fully_convolutional_network/">Image Segmentation using Composable Fully-Convolutional Networks</a> <a class="nav-sublink2" href="/examples/vision/retinanet/">Object Detection with RetinaNet</a> <a class="nav-sublink2" href="/examples/vision/keypoint_detection/">Keypoint Detection with Transfer Learning</a> <a class="nav-sublink2" href="/examples/vision/object_detection_using_vision_transformer/">Object detection with Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/3D_image_classification/">3D image classification from CT scans</a> <a class="nav-sublink2" href="/examples/vision/depth_estimation/">Monocular depth estimation</a> <a class="nav-sublink2" href="/examples/vision/nerf/">3D volumetric rendering with NeRF</a> <a class="nav-sublink2" href="/examples/vision/pointnet_segmentation/">Point cloud segmentation with PointNet</a> <a class="nav-sublink2" href="/examples/vision/pointnet/">Point cloud classification</a> <a class="nav-sublink2" href="/examples/vision/captcha_ocr/">OCR model for reading Captchas</a> <a class="nav-sublink2" href="/examples/vision/handwriting_recognition/">Handwriting recognition</a> <a class="nav-sublink2" href="/examples/vision/autoencoder/">Convolutional autoencoder for image denoising</a> <a class="nav-sublink2" href="/examples/vision/mirnet/">Low-light image enhancement using MIRNet</a> <a class="nav-sublink2" href="/examples/vision/super_resolution_sub_pixel/">Image Super-Resolution using an Efficient Sub-Pixel CNN</a> <a class="nav-sublink2" href="/examples/vision/edsr/">Enhanced Deep Residual Networks for single-image super-resolution</a> <a class="nav-sublink2" href="/examples/vision/zero_dce/">Zero-DCE for low-light image enhancement</a> <a class="nav-sublink2" href="/examples/vision/cutmix/">CutMix data augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/mixup/">MixUp augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/randaugment/">RandAugment for Image Classification for Improved Robustness</a> <a class="nav-sublink2" href="/examples/vision/image_captioning/">Image captioning</a> <a class="nav-sublink2" href="/examples/vision/nl_image_search/">Natural language image search with a Dual Encoder</a> <a class="nav-sublink2" href="/examples/vision/visualizing_what_convnets_learn/">Visualizing what convnets learn</a> <a class="nav-sublink2" href="/examples/vision/integrated_gradients/">Model interpretability with Integrated Gradients</a> <a class="nav-sublink2" href="/examples/vision/probing_vits/">Investigating Vision Transformer representations</a> <a class="nav-sublink2" href="/examples/vision/grad_cam/">Grad-CAM class activation visualization</a> <a class="nav-sublink2" href="/examples/vision/near_dup_search/">Near-duplicate image search</a> <a class="nav-sublink2" href="/examples/vision/semantic_image_clustering/">Semantic Image Clustering</a> <a class="nav-sublink2" href="/examples/vision/siamese_contrastive/">Image similarity estimation using a Siamese Network with a contrastive loss</a> <a class="nav-sublink2" href="/examples/vision/siamese_network/">Image similarity estimation using a Siamese Network with a triplet loss</a> <a class="nav-sublink2 active" href="/examples/vision/metric_learning/">Metric learning for image similarity search</a> <a class="nav-sublink2" href="/examples/vision/metric_learning_tf_similarity/">Metric learning for image similarity search using TensorFlow Similarity</a> <a class="nav-sublink2" href="/examples/vision/nnclr/">Self-supervised contrastive learning with NNCLR</a> <a class="nav-sublink2" href="/examples/vision/video_classification/">Video Classification with a CNN-RNN Architecture</a> <a class="nav-sublink2" href="/examples/vision/conv_lstm/">Next-Frame Video Prediction with Convolutional LSTMs</a> <a class="nav-sublink2" href="/examples/vision/video_transformers/">Video Classification with Transformers</a> <a class="nav-sublink2" href="/examples/vision/vivit/">Video Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/bit/">Image Classification using BigTransfer (BiT)</a> <a class="nav-sublink2" href="/examples/vision/gradient_centralization/">Gradient Centralization for Better Training Performance</a> <a class="nav-sublink2" href="/examples/vision/token_learner/">Learning to tokenize in Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/knowledge_distillation/">Knowledge Distillation</a> <a class="nav-sublink2" href="/examples/vision/fixres/">FixRes: Fixing train-test resolution discrepancy</a> <a class="nav-sublink2" href="/examples/vision/cait/">Class Attention Image Transformers with LayerScale</a> <a class="nav-sublink2" href="/examples/vision/patch_convnet/">Augmenting convnets with aggregated attention</a> <a class="nav-sublink2" href="/examples/vision/learnable_resizer/">Learning to Resize</a> <a class="nav-sublink2" href="/examples/vision/adamatch/">Semi-supervision and domain adaptation with AdaMatch</a> <a class="nav-sublink2" href="/examples/vision/barlow_twins/">Barlow Twins for Contrastive SSL</a> <a class="nav-sublink2" href="/examples/vision/consistency_training/">Consistency training with supervision</a> <a class="nav-sublink2" href="/examples/vision/deit/">Distilling Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/focal_modulation_network/">Focal Modulation: A replacement for Self-Attention</a> <a class="nav-sublink2" href="/examples/vision/forwardforward/">Using the Forward-Forward Algorithm for Image Classification</a> <a class="nav-sublink2" href="/examples/vision/masked_image_modeling/">Masked image modeling with Autoencoders</a> <a class="nav-sublink2" href="/examples/vision/sam/">Segment Anything Model with 🤗Transformers</a> <a class="nav-sublink2" href="/examples/vision/segformer/">Semantic segmentation with SegFormer and Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/vision/simsiam/">Self-supervised contrastive learning with SimSiam</a> <a class="nav-sublink2" href="/examples/vision/supervised-contrastive-learning/">Supervised Contrastive Learning</a> <a class="nav-sublink2" href="/examples/vision/yolov8/">Efficient Object Detection with YOLOV8 and KerasCV</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/vision/'>Computer Vision</a> / Metric learning for image similarity search </div> <div class='k-content'> <h1 id="metric-learning-for-image-similarity-search">Metric learning for image similarity search</h1> Author: <a href="https://twitter.com/mat_kelcey">Mat Kelcey</a> Date created: 2020/06/05 Last modified: 2020/06/09 Description: Example of using similarity metric learning on CIFAR-10 images. <div class='example_version_banner keras_3'>ⓘ This example uses Keras 3</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/vision/ipynb/metric_learning.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/vision/metric_learning.py">GitHub source</a> <hr /> <h2 id="overview">Overview</h2> Metric learning aims to train models that can embed inputs into a high-dimensional space such that "similar" inputs, as defined by the training scheme, are located close to each other. These models once trained can produce embeddings for downstream systems where such similarity is useful; examples include as a ranking signal for search or as a form of pretrained embedding model for another supervised problem. For a more detailed overview of metric learning see: <ul> <li><a href="http://contrib.scikit-learn.org/metric-learn/introduction.html">What is metric learning?</a></li> <li><a href="https://www.youtube.com/watch?v=Jb4Ewl5RzkI">"Using crossentropy for metric learning" tutorial</a></li> </ul> <hr /> <h2 id="setup">Setup</h2> Set Keras backend to tensorflow. <div class="codehilite"><pre><code>import os os.environ["KERAS_BACKEND"] = "tensorflow" import random import matplotlib.pyplot as plt import numpy as np import tensorflow as tf from collections import defaultdict from PIL import Image from sklearn.metrics import ConfusionMatrixDisplay import keras from keras import layers </code></pre></div> <hr /> <h2 id="dataset">Dataset</h2> For this example we will be using the <a href="https://www.cs.toronto.edu/~kriz/cifar.html">CIFAR-10</a> dataset. <div class="codehilite"><pre><code>from keras.datasets import cifar10 (x_train, y_train), (x_test, y_test) = cifar10.load_data() x_train = x_train.astype("float32") / 255.0 y_train = np.squeeze(y_train) x_test = x_test.astype("float32") / 255.0 y_test = np.squeeze(y_test) </code></pre></div> To get a sense of the dataset we can visualise a grid of 25 random examples. <div class="codehilite"><pre><code>height_width = 32 def show_collage(examples): box_size = height_width + 2 num_rows, num_cols = examples.shape[:2] collage = Image.new( mode="RGB", size=(num_cols * box_size, num_rows * box_size), color=(250, 250, 250), ) for row_idx in range(num_rows): for col_idx in range(num_cols): array = (np.array(examples[row_idx, col_idx]) * 255).astype(np.uint8) collage.paste( Image.fromarray(array), (col_idx * box_size, row_idx * box_size) ) # Double size for visualisation. collage = collage.resize((2 * num_cols * box_size, 2 * num_rows * box_size)) return collage # Show a collage of 5x5 random images. sample_idxs = np.random.randint(0, 50000, size=(5, 5)) examples = x_train[sample_idxs] show_collage(examples) </code></pre></div> <img alt="png" src="/img/examples/vision/metric_learning/metric_learning_7_0.png" /> Metric learning provides training data not as explicit <code>(X, y)</code> pairs but instead uses multiple instances that are related in the way we want to express similarity. In our example we will use instances of the same class to represent similarity; a single training instance will not be one image, but a pair of images of the same class. When referring to the images in this pair we'll use the common metric learning names of the <code>anchor</code> (a randomly chosen image) and the <code>positive</code> (another randomly chosen image of the same class). To facilitate this we need to build a form of lookup that maps from classes to the instances of that class. When generating data for training we will sample from this lookup. <div class="codehilite"><pre><code>class_idx_to_train_idxs = defaultdict(list) for y_train_idx, y in enumerate(y_train): class_idx_to_train_idxs[y].append(y_train_idx) class_idx_to_test_idxs = defaultdict(list) for y_test_idx, y in enumerate(y_test): class_idx_to_test_idxs[y].append(y_test_idx) </code></pre></div> For this example we are using the simplest approach to training; a batch will consist of <code>(anchor, positive)</code> pairs spread across the classes. The goal of learning will be to move the anchor and positive pairs closer together and further away from other instances in the batch. In this case the batch size will be dictated by the number of classes; for CIFAR-10 this is 10. <div class="codehilite"><pre><code>num_classes = 10 class AnchorPositivePairs(keras.utils.Sequence): def __init__(self, num_batches): super().__init__() self.num_batches = num_batches def __len__(self): return self.num_batches def __getitem__(self, _idx): x = np.empty((2, num_classes, height_width, height_width, 3), dtype=np.float32) for class_idx in range(num_classes): examples_for_class = class_idx_to_train_idxs[class_idx] anchor_idx = random.choice(examples_for_class) positive_idx = random.choice(examples_for_class) while positive_idx == anchor_idx: positive_idx = random.choice(examples_for_class) x[0, class_idx] = x_train[anchor_idx] x[1, class_idx] = x_train[positive_idx] return x </code></pre></div> We can visualise a batch in another collage. The top row shows randomly chosen anchors from the 10 classes, the bottom row shows the corresponding 10 positives. <div class="codehilite"><pre><code>examples = next(iter(AnchorPositivePairs(num_batches=1))) show_collage(examples) </code></pre></div> <img alt="png" src="/img/examples/vision/metric_learning/metric_learning_13_0.png" /> <hr /> <h2 id="embedding-model">Embedding model</h2> We define a custom model with a <code>train_step</code> that first embeds both anchors and positives and then uses their pairwise dot products as logits for a softmax. <div class="codehilite"><pre><code>class EmbeddingModel(keras.Model): def train_step(self, data): # Note: Workaround for open issue, to be removed. if isinstance(data, tuple): data = data[0] anchors, positives = data[0], data[1] with tf.GradientTape() as tape: # Run both anchors and positives through model. anchor_embeddings = self(anchors, training=True) positive_embeddings = self(positives, training=True) # Calculate cosine similarity between anchors and positives. As they have # been normalised this is just the pair wise dot products. similarities = keras.ops.einsum( "ae,pe->ap", anchor_embeddings, positive_embeddings ) # Since we intend to use these as logits we scale them by a temperature. # This value would normally be chosen as a hyper parameter. temperature = 0.2 similarities /= temperature # We use these similarities as logits for a softmax. The labels for # this call are just the sequence [0, 1, 2, ..., num_classes] since we # want the main diagonal values, which correspond to the anchor/positive # pairs, to be high. This loss will move embeddings for the # anchor/positive pairs together and move all other pairs apart. sparse_labels = keras.ops.arange(num_classes) loss = self.compute_loss(y=sparse_labels, y_pred=similarities) # Calculate gradients and apply via optimizer. gradients = tape.gradient(loss, self.trainable_variables) self.optimizer.apply_gradients(zip(gradients, self.trainable_variables)) # Update and return metrics (specifically the one for the loss value). for metric in self.metrics: # Calling `self.compile` will by default add a [`keras.metrics.Mean`](/api/metrics/metrics_wrappers#mean-class) loss if metric.name == "loss": metric.update_state(loss) else: metric.update_state(sparse_labels, similarities) return {m.name: m.result() for m in self.metrics} </code></pre></div> Next we describe the architecture that maps from an image to an embedding. This model simply consists of a sequence of 2d convolutions followed by global pooling with a final linear projection to an embedding space. As is common in metric learning we normalise the embeddings so that we can use simple dot products to measure similarity. For simplicity this model is intentionally small. <div class="codehilite"><pre><code>inputs = layers.Input(shape=(height_width, height_width, 3)) x = layers.Conv2D(filters=32, kernel_size=3, strides=2, activation="relu")(inputs) x = layers.Conv2D(filters=64, kernel_size=3, strides=2, activation="relu")(x) x = layers.Conv2D(filters=128, kernel_size=3, strides=2, activation="relu")(x) x = layers.GlobalAveragePooling2D()(x) embeddings = layers.Dense(units=8, activation=None)(x) embeddings = layers.UnitNormalization()(embeddings) model = EmbeddingModel(inputs, embeddings) </code></pre></div> Finally we run the training. On a Google Colab GPU instance this takes about a minute. <div class="codehilite"><pre><code>model.compile( optimizer=keras.optimizers.Adam(learning_rate=1e-3), loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), ) history = model.fit(AnchorPositivePairs(num_batches=1000), epochs=20) plt.plot(history.history["loss"]) plt.show() </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/20 77/1000 ━[37m━━━━━━━━━━━━━━━━━━━ 1s 2ms/step - loss: 2.2962 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1700589927.295343 3724442 device_compiler.h:187] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 6s 2ms/step - loss: 2.2504 Epoch 2/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 2.1068 Epoch 3/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 2.0646 Epoch 4/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 2.0210 Epoch 5/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.9857 Epoch 6/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.9543 Epoch 7/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.9175 Epoch 8/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.8740 Epoch 9/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.8474 Epoch 10/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.8380 Epoch 11/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.8146 Epoch 12/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7658 Epoch 13/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7512 Epoch 14/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7671 Epoch 15/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7245 Epoch 16/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7001 Epoch 17/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.7099 Epoch 18/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.6775 Epoch 19/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.6547 Epoch 20/20 1000/1000 ━━━━━━━━━━━━━━━━━━━━ 2s 2ms/step - loss: 1.6356 </code></pre></div> </div> <img alt="png" src="/img/examples/vision/metric_learning/metric_learning_19_3.png" /> <hr /> <h2 id="testing">Testing</h2> We can review the quality of this model by applying it to the test set and considering near neighbours in the embedding space. First we embed the test set and calculate all near neighbours. Recall that since the embeddings are unit length we can calculate cosine similarity via dot products. <div class="codehilite"><pre><code>near_neighbours_per_example = 10 embeddings = model.predict(x_test) gram_matrix = np.einsum("ae,be->ab", embeddings, embeddings) near_neighbours = np.argsort(gram_matrix.T)[:, -(near_neighbours_per_example + 1) :] </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 313/313 ━━━━━━━━━━━━━━━━━━━━ 1s 3ms/step </code></pre></div> </div> As a visual check of these embeddings we can build a collage of the near neighbours for 5 random examples. The first column of the image below is a randomly selected image, the following 10 columns show the nearest neighbours in order of similarity. <div class="codehilite"><pre><code>num_collage_examples = 5 examples = np.empty( ( num_collage_examples, near_neighbours_per_example + 1, height_width, height_width, 3, ), dtype=np.float32, ) for row_idx in range(num_collage_examples): examples[row_idx, 0] = x_test[row_idx] anchor_near_neighbours = reversed(near_neighbours[row_idx][:-1]) for col_idx, nn_idx in enumerate(anchor_near_neighbours): examples[row_idx, col_idx + 1] = x_test[nn_idx] show_collage(examples) </code></pre></div> <img alt="png" src="/img/examples/vision/metric_learning/metric_learning_23_0.png" /> We can also get a quantified view of the performance by considering the correctness of near neighbours in terms of a confusion matrix. Let us sample 10 examples from each of the 10 classes and consider their near neighbours as a form of prediction; that is, does the example and its near neighbours share the same class? We observe that each animal class does generally well, and is confused the most with the other animal classes. The vehicle classes follow the same pattern. <div class="codehilite"><pre><code>confusion_matrix = np.zeros((num_classes, num_classes)) # For each class. for class_idx in range(num_classes): # Consider 10 examples. example_idxs = class_idx_to_test_idxs[class_idx][:10] for y_test_idx in example_idxs: # And count the classes of its near neighbours. for nn_idx in near_neighbours[y_test_idx][:-1]: nn_class_idx = y_test[nn_idx] confusion_matrix[class_idx, nn_class_idx] += 1 # Display a confusion matrix. labels = [ "Airplane", "Automobile", "Bird", "Cat", "Deer", "Dog", "Frog", "Horse", "Ship", "Truck", ] disp = ConfusionMatrixDisplay(confusion_matrix=confusion_matrix, display_labels=labels) disp.plot(include_values=True, cmap="viridis", ax=None, xticks_rotation="vertical") plt.show() </code></pre></div> <img alt="png" src="/img/examples/vision/metric_learning/metric_learning_25_0.png" /> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#metric-learning-for-image-similarity-search'>Metric learning for image similarity search</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#overview'>Overview</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#dataset'>Dataset</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#embedding-model'>Embedding model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#testing'>Testing</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>

CINXE.COM

Metric learning for image similarity search