pytorch-metric-learning

Version 3.0.1

Version 3.0.1 fixes a large time and memory regression in low-precision Euclidean distances, including the default distance used by Contrastive and TripletMargin losses. Public signatures and minimum supported Python, PyTorch, and NumPy versions are unchanged. No migration is required.

Low-precision Euclidean distances

In v3.0.0, float16 and BF16 matrices always used the direct torch.cdist kernel to preserve nearby-point accuracy. Its CUDA backward pass can allocate memory proportional to queries × references × embedding dimensions.

Large eager Euclidean calls now use float32 matrix evaluation when a numerical screen shows that it is safe. Self comparisons reuse prepared inputs and keep their diagonal exactly zero. Nearby or duplicate points, other norms, small matrices, compiled execution, CUDA graph capture, and reduced float32 matmul precision retain the accurate direct calculation. Global precision settings are unchanged.

The improvement is measured on standard-normal inputs; clustered training embeddings can still take the slower direct path. Values and gradients are checked for random, nearby, duplicate, and mixed-scale float16/BF16 inputs.

Measured impact

The fix was measured at commit 32bcfa8, whose numerical code is included unchanged in this release. Each successful timing configuration averages 900 forward/backward pairs across three fresh processes after warmup.

At BF16 batch 2,048, 128 dimensions, and four samples per class:

These are loss/miner measurements, excluding an encoder, optimizer, and data loading. Peak allocation is total PyTorch tensor allocation, not total device usage. The fix does not beat v2.9.0 in every case: RTX 5090 batch-hard triplet still takes 25.6% longer in this configuration. See the complete comparison and limitations for process ranges, controls, numerical checks, and reproduction instructions.

Broader benchmark coverage

The new five-GPU benchmark report compares released v3.0.0 with v2.9.0 on RTX 4090, RTX 5090, L40S, A100, and H100 systems. It covers seven common losses plus batch-hard triplet mining, absolute forward and backward time, total and incremental peak memory, and verified batch limits. CSV tables, SVG scaling charts, raw measurements, and hardware/source provenance are included. Those original release results remain separate from the two-GPU fix measurements.

3.0.0

Version 3 focuses the library on composable metric-learning components used from ordinary PyTorch code. It removes training and evaluation orchestration, makes tensor and extension contracts explicit and typed, and corrects numerical, distributed, and retrieval behavior. It adds losses, evaluation tools, and memory controls.

It is a clean break from version 2, with no compatibility aliases. Look up any version 2 name in the migration guide. This page summarizes the changes since 2.9.0, and each release candidate's notes hold the complete list: rc1, rc2, rc3, rc4, and rc5.

Requirements

• Python 3.12 or newer, with PyTorch 2.14.0 and NumPy 2.5.3 as tested minimums.
• Optional extras: faiss for the Faiss retrieval backends, and clustering for NMI and AMI. The with-hooks, with-hooks-cpu, docs, and dev extras are removed.
• The package is typed and ships py.typed. pytorch_metric_learning.__version__ is removed; use importlib.metadata.version("pytorch-metric-learning").

Breaking changes

• Removed: trainers, testers, logging presets, InferenceModel, the datasets package, and recorded statistics. The training and evaluation recipes show the equivalent PyTorch code.
• Package layout: retrieval lives in evaluation, distributed wrappers in distributed, and utils.indexing replaces utils.loss_and_miner_utils. Every other part of utils is private. See the public API.
• Tensor contracts: ref_emb=None means a query set compared with itself, and passing ref_emb describes a separate set. Miners return named tuples, labels must be on the embeddings' device, and modules are never moved or cast during forward.
• Extension: losses, miners, and regularizers extend torch.nn.Module directly and take keyword-only options. Custom losses implement compute_loss, custom reducers implement element_reduction, and classification losses derive from BaseProxyLoss.
• Arguments: many options are renamed, and the migration guide lists each one.
• Samplers: MPerClassSampler yields whole batches, so pass it as DataLoader(dataset, batch_sampler=sampler). It takes batch_size and num_batches. Seeded sequences differ from version 2.

Changed defaults

• DotProductSimilarity() computes a plain dot product. Pass normalize_embeddings=True for the previous behavior.
• TripletMarginLoss defaults to margin=0.2. Pass margin=0.05 for the previous default.
• AccuracyCalculator searches with exact PyTorch squared Euclidean distance on the input device instead of Faiss.
• Several losses exclude anchors that lack the pairs their term needs, so means average only anchors that have them. Miner weights changed in ProxyAnchorLoss and ProxyNCALoss, and the rc1 notes give the details.

New features

Losses:

• CenterLoss
• CentroidTripletLoss, rewritten for version 3, from an or

3.0.0rc5

This is the fifth release candidate for version 3.0.0. The 3.0.0rc1 notes describe version 3 as a whole, and the migration guide covers version 2 names.

New features

• Compatible losses and miners accept a precomputed distance matrix through the keyword-only mat argument. One matrix can serve a miner and several losses, and a matrix can depend on per-batch metadata that a distance does not receive, without changing the distance object. Losses reject a matrix with no autograd graph, nothing is cached, and unsupported modules raise. See precomputed distance matrices (#211, #697).
• Every miner accepts tuples_per_anchor, which keeps a random subset of each anchor's mined tuples, without replacement. See bounding the output (#193).
• losses.anchor_batched_backward computes the gradient of a tuple loss in chunks of anchors. Its result equals a full-batch backward(), and losses whose objective does not split by anchor are rejected. See anchor_batched_backward (#586).
• losses.DecoupledContrastiveLoss implements DCL and DCLW from Decoupled Contrastive Learning by Yeh et al. See DecoupledContrastiveLoss (#375).
• losses.SelectivelyContrastiveTripletLoss implements the loss from Hard Negative Examples are Hard, but Useful by Xuan et al. See SelectivelyContrastiveTripletLoss (#287).
• regularizers.SphericalEmbeddingRegularizer implements the spherical embedding constraint from Deep Metric Learning with Spherical Embedding by Zhang et al. It differs from CenterInvariantRegularizer, which compares squared norms. See SphericalEmbeddingRegularizer (#346).
• losses.CentroidTripletLoss returns, rewritten for version 3. It implements On the Unreasonable Effectiveness of Centroids in Image Retrieval by Wieczorek et al. Classes of different sizes now work, a class with one embedding provides a negative centroid instead of raising, and swap and triplets_per_anchor are gone. The default distance is now the paper's squared Euclidean distance. Thanks to cwkeam, who contributed the original implementation, and pgrosjean, who diagnosed the failure. See CentroidTripletLoss (#451).
• losses.GE2ELoss implements both variants of Generalized End-to-End Loss for Speaker Verification by Wan et al., with the anchor left out of its own class's

3.0.0rc4

This is the fourth release candidate for version 3.0.0. It adds evaluation utilities and fixed-pair sampling. The 3.0.0rc1 notes describe version 3 as a whole, and the migration guide covers version 2 names.

New features

• evaluation.select_threshold chooses a match/no-match threshold from validation scores, boolean labels, and an explicit metric, and the result works directly with MatchFinder. Built-in metrics are match_accuracy, match_balanced_accuracy, and match_f1. See threshold selection (#733).
• samplers.FixedSetOfPairs draws a fixed, balanced set of positive and negative pairs and yields whole batches, so a pair is never split. The sampler documentation includes a complete DataLoader-to-loss loop (#259).
• evaluation.ndcg_metric registers NDCG from a graded relevance function, normalized by the ideal ranking over the whole reference set. Multi-column labels can drive NDCG and binary metrics in one calculator, and each keeps its own rule for which queries count. See graded relevance and NDCG (#339).

Evaluation changes

• RetrievalInputs gains the required fields reference_labels and ref_includes_query. Only code that constructs RetrievalInputs directly needs to change.
• RetrievalMetric gains exclude_lone_queries, default True. A metric that sets it to False decides through NaN which queries count.

Fixes

• AccuracyCalculator no longer raises TypeError with MPS inputs. The search stays on MPS and metric arithmetic runs on CPU, because MPS has no float64.

Documentation

• The evaluation recipe adds nearest-centroid accuracy, composed from AccuracyCalculator (#205).
• Measuring performance adds benchmark cases that compare separate distance computation with one reused matrix.