# PyTorch Metric Learning

> Losses, miners and testers for training embedding models in PyTorch

- Category: Retrieval & Knowledge (RAG) / Embedding Models & Inference
- Language: Python
- Licence: MIT
- Repository: https://github.com/KevinMusgrave/pytorch-metric-learning
- Website: https://kevinmusgrave.github.io/pytorch-metric-learning/
- Canonical page: https://ai-tldr.dev/tools/pytorch-metric-learning/

## In short

PyTorch Metric Learning is an MIT-licensed Python library for training models whose output is an embedding — a vector where similar items land close together and different ones far apart. It sits in the Embedding Models & Inference category of the Retrieval & Knowledge (RAG) tools and is open-source, written in Python and licensed MIT.

## Overview

PyTorch Metric Learning is an MIT-licensed Python library for training models whose output is an embedding — a vector where similar items land close together and different ones far apart. That is the training side of image retrieval, face and person re-identification, deduplication and contrastive or self-supervised representation learning. The library began during its author Kevin Musgrave's internship at Facebook AI and is described in a 2020 arXiv paper by Musgrave, Serge Belongie and Ser-Nam Lim.

It is organised as nine modules that work on their own or together: losses, miners, distances, reducers, regularizers, samplers, trainers, testers and an AccuracyCalculator. The core pattern drops into an ordinary PyTorch training loop: compute embeddings with your model, pass them with their labels to a loss such as TripletMarginLoss, and call backward. A miner can find the hardest positive and negative pairs in the batch first, and the library converts between pairs and triplets automatically when a loss needs the other form.

Losses are built from swappable parts. A distance object (Euclidean, or cosine similarity) computes the distance matrix, a regularizer adds per-embedding penalties, and a reducer decides how per-pair and per-element losses become one number — for example by keeping only losses above a threshold. A SelfSupervisedLoss wrapper turns any tuple loss into a self-supervised one by comparing embeddings with their augmented views, and CrossBatchMemory supports MoCo-style queues. Release 2.9.0, published to PyPI on 17 August 2025, added SmoothAPLoss; a Datasets module downloads common benchmarks such as CUB-200, Cars196, iNaturalist 2018 and Stanford Online Products.

## What it does

- A large set of metric-learning losses, from TripletMarginLoss and contrastive losses to SupConLoss, CircleLoss, ArcFace variants and SmoothAPLoss
- Miners such as MultiSimilarityMiner and BatchEasyHardMiner that select hard pairs or triplets inside each batch
- Composable losses: swap the distance, add an embedding regularizer, or change the reducer without writing a new loss
- SelfSupervisedLoss wrapper and CrossBatchMemory for self-supervised and MoCo-style training
- Trainers, testers and an AccuracyCalculator (k-NN and clustering metrics, mean reciprocal rank) for complete train/evaluate workflows
- A Datasets module for CUB-200-2011, Cars196, iNaturalist 2018 and Stanford Online Products, plus Colab example notebooks

## Getting started

Versions from 0.9.90 on require torch 1.6 or newer; the other dependencies are numpy, scikit-learn, tqdm and torchvision.

1. Install

   From PyPI, or from conda-forge with `conda install -c conda-forge pytorch-metric-learning`. The `[with-hooks]` extra (or `[with-hooks-cpu]`) adds faiss, record-keeper and tensorboard for the testing and logging modules.

```bash
pip install pytorch-metric-learning
```

2. Add a loss to your training loop

   Embeddings have shape (N, embedding_size) and labels shape (N). TripletMarginLoss builds every valid triplet in the batch from the labels.

```python
from pytorch_metric_learning import losses
loss_func = losses.TripletMarginLoss()

for i, (data, labels) in enumerate(dataloader):
    optimizer.zero_grad()
    embeddings = model(data)
    loss = loss_func(embeddings, labels)
    loss.backward()
    optimizer.step()
```

3. Mine hard pairs

   A miner picks the pairs it judges most difficult and passes them to the loss.

```python
from pytorch_metric_learning import miners, losses
miner = miners.MultiSimilarityMiner()
loss_func = losses.TripletMarginLoss()

hard_pairs = miner(embeddings, labels)
loss = loss_func(embeddings, labels, hard_pairs)
```

4. Customise the loss

   Use cosine similarity instead of Euclidean distance, discard triplet losses above 0.3 and L2-regularise the embeddings.

```python
from pytorch_metric_learning.distances import CosineSimilarity
from pytorch_metric_learning.reducers import ThresholdReducer
from pytorch_metric_learning.regularizers import LpRegularizer
from pytorch_metric_learning import losses

loss_func = losses.TripletMarginLoss(
    distance=CosineSimilarity(),
    reducer=ThresholdReducer(high=0.3),
    embedding_regularizer=LpRegularizer(),
)
```

5. Go self-supervised

   Wrap a tuple loss and compare each batch with an augmented view of itself — no labels needed.

```python
from pytorch_metric_learning.losses import SelfSupervisedLoss, TripletMarginLoss
loss_func = SelfSupervisedLoss(TripletMarginLoss())

embeddings = model(data)
augmented = model(augment(data))
loss = loss_func(embeddings, augmented)
```

## When to use it

- Training an image-retrieval or product-search model where similar items must be nearest neighbours
- Face, person or vehicle re-identification with ArcFace-style or triplet losses
- Self-supervised pre-training of an encoder from augmented views of unlabelled data
- Comparing metric-learning losses and miners on standard benchmarks with the built-in datasets and AccuracyCalculator