Registry / ai-ml / pytorch-metric-learning

pytorch-metric-learning

JSON →
library2.9.0pypypi✓ verified 27d ago

PyTorch Metric Learning is a Python library (version 2.9.0) that simplifies the use of deep metric learning in applications. It offers a modular, flexible, and extensible framework built on PyTorch, providing a wide array of loss functions, miners, samplers, trainers, and testers. The library maintains an active release cadence, with frequent updates introducing new features and improvements.

pip install pytorch-metric-learning
INSTALL
IMPORT
SIG · PYTORCH-METRIC-LEA
P
pytorch-metric-learning
ai-mlpythonv2.9.0
Install
79.8s avg
Import
9695ms
Disk
6554MB
Pass rate
3/ 10
Env Coverage3 / 10
glibc
3.9–3.13
musl
3.9–3.13
Install & Compatibility
Where this runs
tested against v2.9.0 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
musl
glibc
py 3.10
✕ build_error
1/2 runs
py 3.11
✕ build_error
✓ 86.75s
py 3.12
✕ build_error
✓ 79.1s
py 3.13
✕ build_error
✓ 73.5s
py 3.9
✕ build_error
✕ timeout
6554MB installed
● package 6554MB
Code
Verified usage

Verified import paths — ran on the pinned version, not inferred.

TripletMarginLoss
✓ from pytorch_metric_learning.losses import TripletMarginLoss
MultiSimilarityMiner
✓ from pytorch_metric_learning.miners import MultiSimilarityMiner
MetricLossOnly
✓ from pytorch_metric_learning.trainers import MetricLossOnly
✗ from pytorch_metric_learning.trainer import MetricLossOnly
The `trainers` module is plural.
AccuracyCalculator
✓ from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator
DistributedLossWrapper
✓ from pytorch_metric_learning.wrappers import DistributedLossWrapper
✗ from pytorch_metric_learning.losses import DistributedLossWrapper
Wrappers are in their own submodule.

This quickstart demonstrates a basic training loop using PyTorch Metric Learning. It defines a dummy dataset and model, then initializes a `TripletMarginLoss` and `MultiSimilarityMiner`. The training loop moves data and model to the appropriate device, generates embeddings, mines for hard triplets, computes the loss, and performs backpropagation.

import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from pytorch_metric_learning import losses, miners # 1. Dummy Dataset for demonstration class DummyDataset(Dataset): def __init__(self, num_samples=100, embedding_dim=64, num_classes=10): self.data = torch.randn(num_samples, embedding_dim) self.labels = torch.randint(0, num_classes, (num_samples,)) def __len__(self): return len(self.labels) def __getitem__(self, idx): return self.data[idx], self.labels[idx] # 2. Dummy Model (e.g., identity for pre-computed embeddings) class DummyModel(nn.Module): def __init__(self, embedding_dim): super().__init__() self.linear = nn.Linear(embedding_dim, embedding_dim) # Simple linear layer def forward(self, x): return self.linear(x) # Configuration embedding_dim = 64 num_classes = 10 batch_size = 32 num_epochs = 2 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Initialize dataset and dataloader dataset = DummyDataset(embedding_dim=embedding_dim, num_classes=num_classes) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) # Initialize model, loss, and miner model = DummyModel(embedding_dim).to(device) loss_func = losses.TripletMarginLoss(margin=0.1).to(device) miner = miners.MultiSimilarityMiner(epsilon=0.1) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) # Training loop print(f"Training on {device}...") for epoch in range(num_epochs): for i, (data, labels) in enumerate(dataloader): data, labels = data.to(device), labels.to(device) optimizer.zero_grad() embeddings = model(data) # Mine for hard triplets hard_triplets = miner(embeddings, labels) # Compute loss using mined triplets loss = loss_func(embeddings, labels, hard_triplets) loss.backward() optimizer.step() if i % 10 == 0: print(f"Epoch {epoch+1}/{num_epochs}, Batch {i+1}/{len(dataloader)}, Loss: {loss.item():.4f}") print("Training complete.")
Debug
Known issues
breakingThe `emb` argument of `DistributedLossWrapper.forward` was renamed to `embeddings` for consistency with the rest of the library.
fix
Update calls to `DistributedLossWrapper.forward` to use `embeddings` instead of `emb`.
affects: >=2.6.0
breakingThe default value of the `symmetric` flag in `SelfSupervisedLoss` changed from `False` to `True`. If `False`, only `embeddings` are used as anchors. If `True`, `embeddings` and `ref_emb` are both used as anchors.
fix
Explicitly set `symmetric=False` in `SelfSupervisedLoss` initialization if you need the old behavior. Otherwise, be aware of the change in anchor selection.
affects: >=2.2.0
gotchaWhen using `PyTorch's DistributedDataParallel`, `DistributedLossWrapper` and `DistributedMinerWrapper` are essential. Without them, losses and miners in each process will only see a fraction of the global batch, leading to incorrect calculations.
fix
Wrap your loss and miner functions with `DistributedLossWrapper` and `DistributedMinerWrapper` respectively when using `DistributedDataParallel`. Ensure `efficient` parameter is consistent between wrappers if used.
affects: All versions
gotchaVery large batch sizes can lead to `INT_MAX` errors within `loss_and_miner_utils` due to an extremely high number of pairs/triplets being processed.
fix
Reduce your batch size if you encounter `INT_MAX` errors during loss or mining computation.
affects: All versions
gotchaDevice mismatches (CPU/GPU) are a common PyTorch error. Ensure your model, input data, and loss/miner functions are all on the same device (e.g., 'cuda') to avoid `RuntimeError: Expected all tensors to be on the same device...`.
fix
Explicitly move models (`model.to(device)`), data (`data.to(device)`, `labels.to(device)`), and loss/miner functions (`loss_func.to(device)`, `miner.to(device)`) to the target device.
affects: All versions
Errors
Common errors & fixes
ModuleNotFoundError: No module named 'pytorch_metric_learning.loss'
This error occurs when trying to import from `pytorch_metric_learning.loss` (singular) which is an incorrect module path; the correct path uses the plural `losses`.
fix
Use the correct module path: `from pytorch_metric_learning.losses import TripletLoss`
ValueError: Expected `distances` to be non-empty
This error typically occurs when a metric learning loss function, after processing inputs through a miner, receives an empty set of positive or negative pairs/triplets, meaning no valid pairs/triplets were found for the current batch.
fix
Ensure your miner is configured correctly, check if embeddings are discriminative enough, and consider increasing your batch size or adjusting miner margins (e.g., for `TripletMarginMiner`).
AssertionError: The number of embeddings must be equal to the number of labels
The input `embeddings` tensor and `labels` tensor have a different number of samples (i.e., their first dimension sizes do not match), which is a requirement for all metric learning loss functions.
fix
Ensure that `embeddings.shape[0]` (batch size of embeddings) is equal to `labels.shape[0]` (batch size of labels) before passing them to the loss function or trainer.
TypeError: optimizers must be a dictionary or list of dictionaries
When initializing a `pytorch-metric-learning` trainer (e.g., `MetricLossOnly`), the `optimizers` argument was provided as a single `torch.optim.Optimizer` instance instead of the required dictionary mapping model names to optimizers, or a list of such dictionaries.
fix
Pass `optimizers` as a dictionary, e.g., `optimizers = {"model": torch.optim.Adam(model.parameters())}`, where 'model' is the key corresponding to your model in the `models` dictionary passed to the trainer.
Upgrade
Version history
2.9.0latest on PyPI · released Aug 17, 2025
Audit
Dependencies
torchrequiredCore deep learning framework. Requires >= 1.6 for pytorch-metric-learning >= v0.9.90.
numpyrequiredNumerical operations.
scikit-learnrequiredUtility functions, e.g., for data splitting or metrics.
tqdmrequiredProgress bars.
faiss-cpuoptionalEfficient similarity search, part of `[with-hooks]` extra.
faiss-gpuoptionalGPU-accelerated efficient similarity search, part of `[with-hooks]` extra.
record-keeperoptionalLogging and experiment tracking, part of `[with-hooks]` extra.
tensorboardoptionalVisualization tools, part of `[with-hooks]` extra.
Agent activity
18 hits · last 30 days
node
16
OpenAI (training)
1
Resources
pytorch-metric-learning — pip install pytorch-metric-learning · libregistry