Skip to content

3.0.0rc1

This is the first release candidate for version 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 is a clean break: there are no compatibility aliases. Look up any version 2 name in the migration guide.

Requirements

  • Python 3.12 or newer.
  • PyTorch 2.14.0 and NumPy 2.5.3 are the tested minimums; newer versions are permitted.
  • Optional extras: faiss (Faiss 1.15.1 or newer) for the Faiss retrieval backends, and clustering (scikit-learn 1.9.1 or newer) for NMI and AMI.
  • The with-hooks, with-hooks-cpu, docs, and dev extras are removed, so record-keeper and TensorBoard are no longer installed. tqdm is no longer a dependency, and scikit-learn moves to the clustering extra.
  • The package is typed and ships py.typed. Read the installed version with importlib.metadata.version("pytorch-metric-learning"); pytorch_metric_learning.__version__ is removed.

Breaking changes

Package layout

  • Retrieval metrics and backends live in pytorch_metric_learning.evaluation (AccuracyCalculator, CustomKNN, FaissKNN, FaissKMeans, MatchFinder), replacing utils.accuracy_calculator, utils.inference, and utils.retrieval.
  • Distributed wrappers live in pytorch_metric_learning.distributed.
  • utils.loss_and_miner_utils is replaced by utils.indexing. Everything else under utils, including the former common functions, is private.
  • The public API is the __all__ of distances, losses, miners, reducers, regularizers, samplers, evaluation, distributed, types, and utils.indexing; see versioning. Implementation submodules, such as losses.pnp_loss or evaluation.retrieval, are not public import paths, and their helpers are private, including remove_self_comparisons, add_to_index_and_search, and return_results.
  • types exports the aliases and result tuples used in public signatures: PNPVariant, AnchorAggregation, DistributedObjective, SubCenterOutliers, ArrayLike, NeighborLimit, MetricNeighbors, and DatasetLabels.
  • evaluation exports the built-in retrieval metric functions precision_at_k, precision_at_1, r_precision, mean_average_precision, mean_average_precision_at_r, and mean_reciprocal_rank for reuse in custom metrics.

Tensor contracts

  • ref_emb=None means the query set is its own reference set, and self-comparisons are excluded. Passing ref_emb, even the same tensor, describes a separate set. With ref_emb, supply both labels and ref_labels or neither.
  • Labels must be on the embeddings' device; a mismatch raises ValueError. Losses, miners, reducers, and memory never move or cast their parameters and buffers during forward; place modules with .to(device, dtype). Under autocast, parameters keep their configured precision.
  • Miners return types.PairIndices or types.TripletIndices named tuples. Losses, miners, and utils.indexing accept these or plain 3- or 4-tuples.
  • utils.indexing.pairs_to_triplets validates its input like indices_tuple; a triplet tuple raises ValueError.
  • A loss term is {"losses": tensor, "indices": ..., "divisor": optional}. The kind follows from indices: None for an already reduced value, one tensor for per-element losses, an (anchor, other) pair, or a triplet. The reduction_type key is removed.
  • Index tensors in indices_tuple, miner output, and loss terms must be int64 or int32; other dtypes raise TypeError. A loss term's indices is None, one tensor, a pair, or a triplet; 1-tuples are rejected.
  • A supplied indices_tuple with no indices selects nothing. Miner-weighted losses (ArcFace, CosFace, FastAP, LargeMarginSoftmax, NCA, NormalizedSoftmax, ProxyAnchor, ProxyNCA, SmoothAP, SoftTriple, SphereFace, SubCenterArcFace) then return a differentiable zero; pass None for no restriction.
  • Invalid arguments raise TypeError for the wrong type and ValueError for an invalid value, also under python -O. HistogramLoss(num_bins=...), LargeMarginSoftmaxLoss and SphereFaceLoss margins, and SoftTripleLoss or RegularFaceRegularizer with a distance that is not a similarity now raise TypeError.
  • Losses that return one already-reduced value (DynamicSoftMarginLoss, HistogramLoss, ManifoldLoss, P2SGradLoss, PNPLoss, ThresholdConsistentMarginLoss) take no reducer, and P2SGradLoss takes no distance; both were ignored and now raise TypeError. LpRegularizer and ZeroMeanRegularizer likewise take no distance.
  • VICRegLoss requires at least two rows; a one-row batch raises ValueError instead of returning NaN.

Composition and extension

  • BaseMetricLossFunction, BaseMiner, and BaseRegularizer extend torch.nn.Module directly and take keyword-only options typed by types.LossOptions and types.RegularizerOptions. The mixins and composition bases are removed. A reducer is used as given, never copied.
  • Losses with class weights store one row per class or center in a parameter named weight. Except for P2SGradLoss, they derive from the new losses.BaseProxyLoss and take keyword-only weight_regularizer, weight_reg_weight, and weight_init_func.
  • Custom losses implement compute_loss(embeddings, labels, indices_tuple, ref_emb, ref_labels). A custom loss may return {} when nothing is valid; reducers turn an empty dictionary or empty terms into a differentiable zero. sub_loss_names and zero_losses are removed. GenericPairLoss subclasses implement pair_loss(mat, pairs).
  • Custom reducers implement element_reduction(losses, anchors, embeddings, labels); the pair and triplet hooks are removed. Terms with a divisor (ProxyAnchorLoss, MarginLoss, SparseCentersRegularizer) require DivisorReducer, alone or inside MultipleReducers; other reducers raise ValueError instead of ignoring the divisor.
  • Regularization penalties are added after reduction. loss_terms(...) returns the unreduced objective and penalty terms.
  • Miners take forward(embeddings, labels, indices_tuple=None, ref_emb=None, ref_labels=None) and restrict their search to optional candidate tuples. This argument replaces HDCMiner.set_idx_externally and reset_idx. The third positional argument of miners and DistributedMinerWrapper is now indices_tuple, so pass references by keyword.
  • Distance options are keyword-only. BaseDistance.is_inverted is a class attribute. BaseDistance.pairwise(query, ref) computes normalized, powered row-wise values; smallest_dist and largest_dist are removed.
  • DistributedLossWrapper and DistributedMinerWrapper take objective="global" or "local_anchor" instead of efficient.

Renamed or changed arguments

  • ArcFaceLoss and SubCenterArcFaceLoss margins are in radians (default 0.5). LargeMarginSoftmaxLoss requires an integer margin.
  • CircleLoss(margin=...) replaces m; HistogramLoss(num_bins=...) replaces n_bins and delta; RankedListLoss(neg_temperature=..., pos_temperature=...) replaces Tn and Tp; AngularMiner(alpha=...) replaces angle; PerAnchorReducer(aggregation="mean") or aggregation="sum" replaces aggregation_func.
  • ManifoldLoss(embedding_size, num_proxies, lambda_c, ...) replaces l, K, and lambdaC; proxy assignments are passed as labels.
  • P2SGradLoss(num_classes, embedding_size) matches the other classification losses. MultipleLosses takes mappings for losses, miners, and weights, accepts only BaseMetricLossFunction or MultipleLosses values, and has the same forward signature as a loss, including ref_emb and ref_labels. CrossBatchMemory, SelfSupervisedLoss, and DistributedLossWrapper accept it (types.MetricLoss). Its per-loss miners select from a supplied indices_tuple as candidates.
  • CosineSimilarity() takes no options; use DotProductSimilarity(normalize_embeddings=True, ...) for other norms or powers.
  • LiftedStructureLoss and GeneralizedLiftedStructureLoss take (pos_margin, neg_margin) positionally, like ContrastiveLoss.
  • MarginLoss names its main term "loss"; DynamicSoftMarginLoss.hist_ is hist.
  • SubCenterArcFaceLoss.get_outliers(embeddings, labels, threshold_radians=...) takes its threshold in radians and always returns SubCenterOutliers(outliers, dominant_centers).
  • MPerClassSampler(num_samples=...) replaces length_before_new_iter.
  • MatchFinder.get_matching_pairs is replaced by get_match_matrix (a boolean matrix) and get_match_indices (an (num_matches, 2) index tensor).
  • AccuracyCalculator no longer takes return_per_class. get_accuracy returns dict[str, float]; get_accuracy_per_class returns values keyed by class label, with NaN for classes without relevant references. Custom metrics are passed as custom_metrics={name: RetrievalMetric(function, neighbors=...)} or ClusteringMetric(function); a retrieval metric returns one value per query and the calculator owns all reduction. metric_functions(), requires_knn(), requires_clustering(), aggregate(), and the calculate_* methods are removed. Naming a clustering metric in get_accuracy_per_class(include=...) raises ValueError.
  • Samplers draw from PyTorch generators and accept a keyword-only generator; seeded sequences differ from version 2.

Removed

  • Trainers, testers, logging presets, InferenceModel, run_pca, and the datasets package. The training and evaluation recipes show the equivalent PyTorch code, including classifier, proxy, two-stream, cascaded, adversarial, and memory-warmup training, split composition, and parent-conditioned metrics.
  • Recorded statistics (collect_stats, COLLECT_STATS, and attributes such as miner.pos_pair_dist). Compute diagnostics from returned tensors or with forward hooks. Recording forced host synchronization inside forward.
  • DoNothingReducer (use loss_terms), EmbeddingsAlreadyPackagedAsTriplets, TuplesToWeightsSampler, BaseLossWrapper, the common_functions logging, tensor-conversion, checkpoint, and label-mapping helpers, and Jupyter notebooks.

Changed defaults

  • AccuracyCalculator searches with exact PyTorch squared Euclidean distance (CustomKNN) on the input device instead of Faiss, and selects a device only when given one. NMI and AMI are computed only when included with a kmeans_func.
  • A fixed k is capped at the neighbors available in each call.
  • R-precision with an integer k below a query's relevant count divides by the relevant count instead of by k. That configuration still logs a warning; use k="max_bin_count".
  • utils.indexing.convert_to_triplets defaults to all triplets.
  • ProxyAnchorLoss applies miner weights as multiplicative weights on the exponential terms; version 2 added w - 1 to the exponent. Results without a miner are unchanged.
  • ProxyNCALoss miner weights count every position in the mined tuples, like the other classification losses; version 2 counted only anchors. A mined triplet now also weights its positive and negative rows.
  • PerAnchorReducer(aggregation="mean") averages over anchors that have losses; version 2 averaged over every row of the batch.
  • CircleLoss, GeneralizedLiftedStructureLoss, MultiSimilarityLoss, SupConLoss, SmoothAPLoss, and PNPLoss exclude anchors that lack the pairs their term needs, so means average only anchors that have them, as NCALoss and FastAPLoss already did.
  • SupConLoss follows the paper when a batch has positives but no negatives, instead of returning zero.

Fixed

  • NCALoss and ProxyNCALoss compute the objective in log space. A single query with valid separate references no longer returns zero, and large softmax_scale values no longer underflow hard examples to zero loss and gradient. Queries without a positive reference are excluded.
  • NTXentLoss computes per-anchor log-sum-exp directly; float16 hard positives keep their loss and gradients.
  • AccuracyCalculator relevance counts, lone-query detection, and k="max_bin_count" honor custom label comparators, so a comparator that rejects self-matches no longer produces infinite MAP or NaN R-precision.
  • CrossBatchMemory rejects labels that change value when converted to label_dtype, and validates inputs and runs the wrapped miner and loss before modifying its queue. Its miner no longer sees each anchor's own enqueued copy, which an easy-positive miner used to select and then lose, leaving no positive.
  • PerAnchorReducer aggregates the supplied losses by anchor, supports references larger than the query batch and duplicate pairs, and no longer builds a dense matrix.
  • AngularMiner measures angles on normalized embeddings, matching AngularLoss.
  • LiftedStructureLoss positive pairs without negatives contribute zero, as in the paper; the loss rejects ref_emb, which it mishandled.
  • SmoothAPLoss derives positives from labels instead of positional blocks.
  • ThresholdConsistentMarginLoss uses ref_emb and its configured similarity.
  • TupletMarginLoss no longer returns NaN for identical anchors and positives, and AngularLoss normalizes negatives.
  • RankedListLoss computes exact weighted means without a 1e-5 offset and cannot overflow for large temperatures.
  • TripletMarginLoss(swap=True) measures positive-negative distances on the reference set.
  • FaissKNN uses a GPU only with a GPU build of Faiss.
  • BatchHardMiner, BatchEasyHardMiner, DistanceWeightedMiner, and MultiSimilarityMiner return empty tuples for an empty query or reference set instead of raising IndexError.
  • AccuracyCalculator runs under torch.no_grad() and builds no autograd graph. A knn_func returning indices outside the reference set, such as a FaissKNN(reset_before=False) whose index has grown, raises ValueError instead of IndexError, and so do indices not shaped (len(query), k).
  • float16 and bfloat16 normalization runs in float32, including in P2SGradLoss, so zero and tiny rows no longer give NaN values or gradients. ThresholdConsistentMarginLoss sums its penalties and DynamicSoftMarginLoss counts its histogram in float32, so large half-precision batches no longer overflow. LpDistance.pairwise no longer adds 1e-6, so identical rows give 0, matching the distance matrix.
  • Distributed wrappers raise invalid-input errors on every rank instead of leaving other ranks waiting in a collective.
  • AngularLoss takes negatives from labels even when a miner supplies only positives, instead of returning zero.
  • GeneralizedLiftedStructureLoss no longer adds a hinge for anchors missing positives or negatives.
  • PNPLoss, labeled ManifoldLoss, and P2SGradLoss give a differentiable zero for empty batches instead of NaN.
  • Miner weights normalize occurrence counts before casting, so float16 counts above 65504 no longer become infinite.
  • DistanceWeightedMiner normalizes weights per anchor over eligible negatives, so default cutoffs no longer underflow to an empty output.
  • LpDistance preserves small float16 and bfloat16 distances and their gradients.
  • Losses never move or cast learnable parameters during forward; one float64 call previously converted a float32 proxy matrix permanently.
  • AccuracyCalculator keeps label dtypes, so large integer identifiers no longer collide (reported and fixed by Akshit Khokhani in #763).
  • MPerClassSampler yields exactly len(sampler) indices.
  • MatchFinder.is_match applies the configured normalization and power.

Performance

  • Pair-to-triplet conversion and all-triplet enumeration group negatives by anchor; memory scales with the emitted triplets instead of a positives-by-negatives boolean matrix.
  • NTXentLoss, TupletMarginLoss, and LiftedStructureLoss reduce negatives per anchor instead of materializing positive-by-negative matrices.
  • Default reducers, self-comparison removal, and random triplet sampling avoid data-dependent host synchronization where the output shape allows it.
  • AccuracyCalculator retrieves only the neighbors the selected metrics need, and counts relevant references from label counts with the default comparator. With only precision_at_1 and k=None, 2000 queries against 20000 references took 0.10 s instead of 2.9 s on one CPU thread.
  • LpDistance computes half-precision matrices with torch.cdist in float32 instead of expanding every row pair.

See performance measurements for the measured workloads.