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:
faissfor the Faiss retrieval backends, andclusteringfor NMI and AMI. Thewith-hooks,with-hooks-cpu,docs, anddevextras are removed. - The package is typed and ships
py.typed.pytorch_metric_learning.__version__is removed; useimportlib.metadata.version("pytorch-metric-learning").
Breaking changes¶
- Removed: trainers, testers, logging presets,
InferenceModel, thedatasetspackage, and recorded statistics. The training and evaluation recipes show the equivalent PyTorch code. - Package layout: retrieval lives in
evaluation, distributed wrappers indistributed, andutils.indexingreplacesutils.loss_and_miner_utils. Every other part ofutilsis private. See the public API. - Tensor contracts:
ref_emb=Nonemeans a query set compared with itself, and passingref_embdescribes a separate set. Miners return named tuples, labels must be on the embeddings' device, and modules are never moved or cast duringforward. - Extension: losses, miners, and regularizers extend
torch.nn.Moduledirectly and take keyword-only options. Custom losses implementcompute_loss, custom reducers implementelement_reduction, and classification losses derive fromBaseProxyLoss. - Arguments: many options are renamed, and the migration guide lists each one.
- Samplers:
MPerClassSampleryields whole batches, so pass it asDataLoader(dataset, batch_sampler=sampler). It takesbatch_sizeandnum_batches. Seeded sequences differ from version 2.
Changed defaults¶
DotProductSimilarity()computes a plain dot product. Passnormalize_embeddings=Truefor the previous behavior.TripletMarginLossdefaults tomargin=0.2. Passmargin=0.05for the previous default.AccuracyCalculatorsearches 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
ProxyAnchorLossandProxyNCALoss, and the rc1 notes give the details.
New features¶
Losses:
- CenterLoss
- CentroidTripletLoss, rewritten for version 3, from an original implementation by cwkeam and a diagnosis by pgrosjean
- DecoupledContrastiveLoss
- DynamicArcFaceLoss, adapted from a contribution by chingis
- GE2ELoss
- MultiSupConLoss, adapted from a contribution by Qibin Liang
- RankNContrastLoss
- SelectivelyContrastiveTripletLoss
Everything else:
regularizers.SphericalEmbeddingRegularizerimplements the spherical embedding constraint.losses.anchor_batched_backwardcomputes a tuple loss's gradient in chunks of anchors, for batches that do not fit in memory.- Compatible losses and miners accept a
precomputed distance matrix through
mat. - Every miner accepts
tuples_per_anchor. SelfSupervisedLossaccepts more than two views.CrossBatchMemoryaccepts label sets, andDistributedLossWrappersupports it withobjective="local_anchor".evaluation.select_thresholdchooses a match threshold from validation scores.evaluation.ndcg_metricadds NDCG from a graded relevance function.samplers.FixedSetOfPairsdraws a fixed, balanced set of pairs.FaissKNNpasses CUDA tensors directly to a GPU index on the same device, extracted from a contribution by Domenico Muscillo.- Recipes cover ProxyNCA++, N-pairs class mining, and nearest-centroid accuracy.
Numerical corrections¶
- Unnormalized distances no longer overflow float16: they are computed and returned in float32. See calling a distance.
NCALoss,ProxyNCALoss,NTXentLoss,ManifoldLoss, andRankedListLosscompute in log space or without offsets, so they no longer underflow, overflow, or return NaN where they did.ManifoldLossraisesTypeErrorfor a distance that is not a similarity.NTXentLosspreserves small losses under large score offsets (contributed by Afloat), andTupletMarginLossuses 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.
CrossBatchMemoryvalidates inputs before changing its queue, andDynamicSoftMarginLosskeeps its histogram detached across optimizer steps.AccuracyCalculatorkeeps 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, andLiftedStructureLossreduce negatives per anchor instead of building positive-by-negative matrices.AccuracyCalculatorretrieves only the neighbors the selected metrics need.
See measuring performance for workloads and commands.
Changes since 3.0.0rc5¶
anchor_batched_backwardkeeps what the input dtype selects in a full loss. Float16 and bfloat16 rows previously got float32's normalizationepsandtorch.cdistkernel, so their loss and gradients could differ materially from a fullbackward(). Chunks also stay out of autocast, whose float16 terms could overflow the running loss toinf.- Installation instructions no longer point at an older release candidate.