Skip to content

Use PyTorch Metric Learning

Choose a loss for embeddings produced by your PyTorch model, then add mining, custom distances, reduction, or regularization as needed. Applications own training loops, data loading, optimizers, and logging.

Install

pip install pytorch-metric-learning
pip install "pytorch-metric-learning[faiss]"       # Faiss retrieval backends
pip install "pytorch-metric-learning[clustering]"  # NMI and AMI metrics

Version 3 requires Python 3.12 or newer, PyTorch 2.14.0 or newer, and NumPy 2.5.3 or newer. See the release notes, the migration guide, and the version 2 documentation. New to the library? Read the concepts overview.

Your PyTorch code computes embeddings and calls the library's miner and loss to get a scalar loss, then calls backward and steps its optimizer. For evaluation, your code extracts embeddings and AccuracyCalculator retrieves nearest neighbors to compute metrics.

Choose by task

Task Public entry point Contract and runnable composition
Learn from labels and pairs losses.ContrastiveLoss Contrastive loss, training recipe
Learn relative distances losses.TripletMarginLoss Triplet margin loss
Train learnable class proxies losses.NormalizedSoftmaxLoss, losses.ProxyAnchorLoss Proxy parameters and optimizer ownership
Select difficult pairs or triplets miners Miners, cascaded mining
Compare separate embedding streams Loss ref_emb / ref_labels arguments Two-stream gradients
Use augmented views losses.SelfSupervisedLoss, losses.VICRegLoss Self-supervision
Reuse detached embeddings losses.CrossBatchMemory Memory contract, warmup
Balance batches by label samplers.MPerClassSampler, samplers.HierarchicalSampler Samplers
Sample from mined tuple frequency utils.indexing.convert_to_weights Weighted sampling
Evaluate retrieval evaluation.AccuracyCalculator Metrics and query/reference rules, evaluation recipe
Evaluate hierarchical or two-stream labels AccuracyCalculator with explicit subsets Parent-conditioned and two-stream evaluation
Find nearest neighbors or threshold matches evaluation.CustomKNN, evaluation.MatchFinder Retrieval backends
Train across devices distributed.DistributedLossWrapper Distributed values and gradients
Generate hard negatives Ordinary PyTorch modules and triplet loss Adversarial objectives

Component references

Execute a recipe

The training and evaluation recipes list their inputs, expected results, supported composition, and exact verification command. They use synthetic tensors and download nothing.

Imports

from pytorch_metric_learning import (
    distances,
    losses,
    miners,
    reducers,
    regularizers,
    samplers,
)
from pytorch_metric_learning.evaluation import (
    AccuracyCalculator,
    CustomKNN,
    MatchFinder,
)
from pytorch_metric_learning.distributed import DistributedLossWrapper
from pytorch_metric_learning.utils import indexing
from pytorch_metric_learning.types import PairIndices, TripletIndices

utils.indexing holds the public pair, triplet, and weight utilities used by custom losses and miners; other modules under utils are private. Shared type annotations live in pytorch_metric_learning.types. Names starting with an underscore are private.