PyTorch Metric Learning

// 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.

pip install pytorch-metric-learning

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()

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)

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(),
)

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

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

Commands and code are distilled from the project's own documentation — always check the official repo for the latest.

// 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

// How PyTorch Metric Learning compares

PyTorch Metric Learning alongside other open-source embedding models & inference tools AI/TLDR tracks, ranked by GitHub stars.