Skip to content

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:

Everything else:

Numerical corrections

  • Unnormalized distances no longer overflow float16: they are computed and returned in float32. See calling a distance.
  • NCALoss, ProxyNCALoss, NTXentLoss, ManifoldLoss, and RankedListLoss compute in log space or without offsets, so they no longer underflow, overflow, or return NaN where they did. ManifoldLoss raises TypeError for a distance that is not a similarity.
  • NTXentLoss preserves small losses under large score offsets (contributed by Afloat), and TupletMarginLoss uses exact endpoint angles (adapted from a contribution by AryanRG).
  • Float16 and bfloat16 normalization and counts run in float32, so zero and tiny rows no longer give NaN and large batches no longer overflow.
  • CrossBatchMemory validates inputs before changing its queue, and DynamicSoftMarginLoss keeps its histogram detached across optimizer steps.
  • AccuracyCalculator keeps label dtypes, so large identifiers no longer collide (reported and fixed by Akshit Khokhani).

Regression tests from kocean9-freedom and Afloat cover DistanceWeightedMiner, PNPLoss, and HistogramLoss.

Performance

  • Pair-to-triplet conversion and all-triplet enumeration scale with the emitted triplets.
  • NTXentLoss, TupletMarginLoss, and LiftedStructureLoss reduce negatives per anchor instead of building positive-by-negative matrices.
  • AccuracyCalculator retrieves only the neighbors the selected metrics need.

See measuring performance for workloads and commands.

Changes since 3.0.0rc5

  • anchor_batched_backward keeps what the input dtype selects in a full loss. Float16 and bfloat16 rows previously got float32's normalization eps and torch.cdist kernel, so their loss and gradients could differ materially from a full backward(). Chunks also stay out of autocast, whose float16 terms could overflow the running loss to inf.
  • Installation instructions no longer point at an older release candidate.