KevinMusgrave/pytorch-metric-learning

Train embedding models with composable PyTorch losses, hard-example miners, distances, reducers, and regularizers. Use them in your own training loop for retrieval, clustering, recognition, or learning from pairs and triplets.

Docs for humans · Documentation source

Install

Install from PyPI with uv or pip:

uv pip install pytorch-metric-learning
pip install pytorch-metric-learning

The core package needs only PyTorch and NumPy. Optional extras add the Faiss retrieval backend, pytorch-metric-learning[faiss], and scikit-learn for the NMI/AMI clustering metrics, pytorch-metric-learning[clustering].

Version 3 requires Python 3.12 or newer, PyTorch 2.14.0 or newer, and NumPy 2.5.3 or newer. Upgrading from version 2? Read the release notes and look up old names in the migration guide. The version 2 documentation remains available.

A minimal training step

Your model produces one embedding per input. Equal labels identify positive pairs; different labels identify negatives. Here a contrastive loss pulls positive pairs together and pushes negatives apart:

import torch
from pytorch_metric_learning import losses

torch.manual_seed(17)
model = torch.nn.Linear(8, 4)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data = torch.randn(12, 8)
labels = torch.arange(4).repeat_interleave(3)
loss_func = losses.ContrastiveLoss(neg_margin=1.4)

optimizer.zero_grad(set_to_none=True)
loss = loss_func(model(data), labels)
loss.backward()
optimizer.step()

The model returns embeddings shaped (12, 4) and labels have shape (12,). This loss L2-normalizes embeddings by default. It averages active positive and negative pair penalties separately, then adds them. See the contrastive-loss contract for distances, margins, and reduction options.

Learn more

The documentation starts with a task index and a concepts overview of embeddings, tuples, distances, loss terms, and reducers. Every component page states its signature, defaults, tensor contract, and paper, and the recipes show complete training and evaluation code.

Contributing

See CONTRIBUTING.md. Thanks to all contributors, upstream projects, and supporters.

Citing this library

If you use this library in your research, please cite the PyTorch Metric Learning paper:

@article{Musgrave2020PyTorchML,
  title={PyTorch Metric Learning},
  author={Kevin Musgrave and Serge J. Belongie and Ser-Nam Lim},
  journal={ArXiv},
  year={2020},
  volume={abs/2008.09164}
}