最新消息

8月17日: v2.9.0

• 新增了SmoothAPLoss。
• 改进了SubCenterArcFaceLoss和GenericPairLoss。
• 感谢ir2718、lucamarini22和marcpaga。

12月11日: v2.8.0

• 新增了Datasets模块，方便下载常用数据集：
  • CUB200
  • Cars196
  • INaturalist 2018
  • Stanford Online Products
• 感谢ir2718。

文档

• 在此查看文档
• 在此查看安装说明
• 在此查看可用的损失函数、挖掘器等

Google Colab示例

请参阅examples文件夹，其中包含可下载或在Google Colab上运行的笔记本。

PyTorch度量学习概述

该库包含9个模块，每个模块都可以独立地集成到您现有的代码库中，也可以组合使用以实现完整的训练/测试流程。

损失函数的工作原理

在训练循环中使用损失函数和挖掘器

让我们初始化一个普通的TripletMarginLoss：

from pytorch_metric_learning import losses
loss_func = losses.TripletMarginLoss()

要在训练循环中计算损失，需要将模型计算出的嵌入向量和对应的标签传入。嵌入向量的形状应为(N, embedding_size)，标签的形状应为(N)，其中N是批次大小。

# 您的训练循环
for i, (data, labels) in enumerate(dataloader):
	optimizer.zero_grad()
	embeddings = model(data)
	loss = loss_func(embeddings, labels)
	loss.backward()
	optimizer.step()

TripletMarginLoss会根据您传入的标签，在批次内计算所有可能的三元组。同标签的嵌入构成锚点-正样本对，不同标签的则构成锚点-负样本对。

有时添加一个挖掘器会有帮助：

from pytorch_metric_learning import miners, losses
miner = miners.MultiSimilarityMiner()
loss_func = losses.TripletMarginLoss()

# 您的训练循环
for i, (data, labels) in enumerate(dataloader):
	optimizer.zero_grad()
	embeddings = model(data)
	hard_pairs = miner(embeddings, labels)
	loss = loss_func(embeddings, labels, hard_pairs)
	loss.backward()
	optimizer.step()

在上述代码中，挖掘器会找到它认为特别困难的正样本和负样本对。请注意，尽管TripletMarginLoss是以三元组为基础的，但也可以传入成对的数据。这是因为该库会在必要时自动将成对数据转换为三元组，或将三元组转换为成对数据。

自定义损失函数

损失函数可以通过距离度量、规约器和正则化器进行定制。下图中，挖掘器找到了批次内困难样本的索引，这些索引用于从距离对象计算出的距离矩阵中提取对应值。对于这张图来说，损失函数是基于成对计算的，因此它会为每一对样本计算损失。此外，还提供了一个正则化器，这样就可以为批次中的每个嵌入计算正则化损失。成对损失和单个元素的损失会被传递给规约器，在这张图中，规约器只保留高值的损失。最后，对高值的成对损失和元素损失分别求平均，并相加得到最终的损失。

下面是一个自定义TripletMarginLoss的例子：

from pytorch_metric_learning.distances import CosineSimilarity
from pytorch_metric-learning.reducers import ThresholdReducer
from pytorch_metric-learning.regularizers import LpRegularizer
from pytorch_metric_learning import losses
loss_func = losses.TripletMarginLoss(distance = CosineSimilarity(), 
				     reducer = ThresholdReducer(high=0.3), 
			 	     embedding_regularizer = LpRegularizer())

这个自定义的三元组损失具有以下特性：

• 损失将使用余弦相似度而非欧氏距离来计算。
• 所有高于0.3的三元组损失都将被丢弃。
• 嵌入将接受L2正则化。

将损失函数用于无监督/自监督学习

为自监督学习提供了一个SelfSupervisedLoss包装器：

from pytorch_metric_learning.losses import SelfSupervisedLoss
loss_func = SelfSupervisedLoss(TripletMarginLoss())

# 您的训练循环
for i, data in enumerate(dataloader):
	optimizer.zero_grad()
	embeddings = your_model(data)
	augmented = your_model(your_augmentation(data))
	loss = loss_func(embeddings, augmented)
	loss.backward()
	optimizer.step()

如果您对MoCo-风格的自监督感兴趣，请查看MoCo on CIFAR10笔记本。它使用CrossBatchMemory来实现动量编码器队列，这意味着您可以使用任何成对损失函数和成对挖掘器从队列中提取困难样本。

库的其余部分亮点

• 如果您想以一种便捷的方式训练模型，请查看训练器。
• 想在数据集上测试模型的准确率吗？试试测试器。
• 要直接计算嵌入空间的准确率，可以使用AccuracyCalculator。

如果您时间有限，想要一个完整的训练/测试流程，可以查看示例 Google Colab 笔记本。

如需了解更多关于上述内容的信息，请参阅文档。

安装

所需的 PyTorch 版本

• pytorch-metric-learning >= v0.9.90 需要 torch >= 1.6
• pytorch-metric-learning < v0.9.90 没有版本要求，但已在 torch >= 1.2 上进行过测试

其他依赖：numpy, scikit-learn, tqdm, torchvision

Pip

pip install pytorch-metric-learning

获取最新开发版本：

pip install pytorch-metric-learning --pre

在 Windows 上安装：

pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install pytorch-metric-learning

安装带有评估和日志记录功能的版本

（这将安装 faiss-gpu 的非官方 pypi 版本，以及 record-keeper 和 tensorboard）：

pip install pytorch-metric-learning[with-hooks]

安装带有评估和日志记录功能的 CPU 版本

（这将安装 faiss-cpu 的非官方 pypi 版本，以及 record-keeper 和 tensorboard）：

pip install pytorch-metric-learning[with-hooks-cpu]

Conda

conda install -c conda-forge pytorch-metric-learning

要使用测试模块，您需要 faiss，也可以通过 conda 安装。请参阅 faiss 的安装说明。

基准测试结果

请参阅 powerful-benchmarker 查看基准测试结果并使用该工具。

开发

开发工作在 dev 分支上进行：

git checkout dev

单元测试可以使用默认的 unittest 库运行：

python -m unittest discover

您可以使用环境变量指定测试的数据类型和设备。例如，要在 CPU 上使用 float32 和 float64 进行测试：

TEST_DTYPES=float32,float64 TEST_DEVICE=cpu python -m unittest discover

如果只想运行单个测试文件而不是整个测试套件，可以指定文件名：

python -m unittest tests/losses/test_angular_loss.py

代码使用 black 和 isort 格式化：

pip install black isort
./format_code.sh

致谢

贡献者

感谢所有提交拉取请求的贡献者！

贡献者	亮点
domenicoMuscill0	- ManifoldLoss - P2SGradLoss - HistogramLoss - DynamicSoftMarginLoss - RankedListLoss
mlopezantequera	- 使测试器能够在任意查询集和参考集的组合上运行 - 使AccuracyCalculator能够处理任意标签比较
cwkeam	- SelfSupervisedLoss - VICRegLoss - 向AccuracyCalculator添加了平均倒数排名准确率 - BaseLossWrapper
ir2718	- ThresholdConsistentMarginLoss - SmoothAPLoss - Datasets模块
marijnl	- BatchEasyHardMiner - TwoStreamMetricLoss - GlobalTwoStreamEmbeddingSpaceTester - 使用trainers.TwoStreamMetricLoss的示例
chingisooinar	SubCenterArcFaceLoss
elias-ramzi	HierarchicalSampler
fjsj	SupConLoss
AlenUbuntu	CircleLoss
interestingzhuo	PNPLoss
wconnell	学习scRNAseq度量嵌入
mkmenta	改进了get_all_triplets_indices（修复了INT_MAX错误）
AlexSchuy	优化了utils.loss_and_miner_utils.get_random_triplet_indices
JohnGiorgi	utils.distributed中的all_gather
Hummer12007	utils.key_checker
vltanh	使InferenceModel.train_indexer能够接受数据集
btseytlin	InferenceModel中的get_nearest_neighbors
mlw214	向AccuracyCalculator添加了return_per_class
layumi	InstanceLoss
NoTody	帮助在分布式包装器中添加了ref_emb和ref_labels。
ElisonSherton	修复了ArcFaceLoss中的一个边缘情况。
stompsjo	改进了NTXentLoss的文档。
Puzer	修复了PNPLoss中的一个错误。
elisim	对DistributedLossWrapper进行了开发者层面的改进。
lucamarini22
marcpaga
GaetanLepage
z1w
thinline72
tpanum
fralik
joaqo
JoOkuma
gkouros
yutanakamura-tky
KinglittleQ
martin0258
michaeldeyzel
HSinger04
rheum
bot66

Facebook AI

感谢Facebook AI的Ser-Nam Lim，以及我的研究导师Serge Belongie教授。这个项目始于我在Facebook AI的实习期间，当时我从Ser-Nam及其计算机视觉和机器学习工程师、研究科学家团队那里获得了宝贵的反馈。特别要感谢Ashish Shah和Austin Reiter，他们在项目早期阶段对我的代码进行了评审。

开源仓库

本库包含从以下优秀开源仓库改编和修改的代码：

• https://github.com/bnu-wangxun/Deep_Metric
• https://github.com/chaoyuaw/incubator-mxnet/blob/master/example/gluon/embedding_learning
• https://github.com/facebookresearch/deepcluster
• https://github.com/geonm/proxy-anchor-loss
• https://github.com/idstcv/SoftTriple
• https://github.com/kunhe/FastAP-metric-learning
• https://github.com/ronekko/deep_metric_learning
• https://github.com/tjddus9597/Proxy-Anchor-CVPR2020
• http://kaizhao.net/regularface
• https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts

标志

感谢Jeff Musgrave设计了标志。

引用本库

如果您希望在论文中引用 pytorch-metric-learning，可以使用以下 BibTeX 格式：

@article{Musgrave2020PyTorchML,
  title={PyTorch Metric Learning},
  author={Kevin Musgrave and Serge J. Belongie and Ser-Nam Lim},
  journal={ArXiv},
  year={2020},
  volume={abs/2008.09164}
}

PyTorch Metric Learning 快速上手指南

pytorch-metric-learning 是一个功能强大的 PyTorch 库，提供了丰富的损失函数、挖掘器（Miners）、距离度量等模块，旨在简化度量学习（Metric Learning）模型的训练与评估流程。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS, 或 Windows
• Python: 建议 Python 3.6+
• PyTorch:
  • 版本 v0.9.90 及以上需要 torch >= 1.6
  • 旧版本已在 torch >= 1.2 上测试通过
• 其他依赖: numpy, scikit-learn, tqdm, torchvision

提示：如果您需要使用评估和日志记录功能（如 FAISS），安装时可选择附加选项（见下文）。

安装步骤

您可以选择通过 pip 或 conda 进行安装。国内开发者建议使用清华或阿里镜像源以加速下载。

方式一：使用 Pip 安装

标准安装：

pip install pytorch-metric-learning -i https://pypi.tuna.tsinghua.edu.cn/simple

获取最新开发版：

pip install pytorch-metric-learning --pre -i https://pypi.tuna.tsinghua.edu.cn/simple

安装包含评估和日志功能的版本（推荐）： 此选项会自动安装 faiss (GPU 或 CPU 版)、record-keeper 和 tensorboard。

• GPU 版本：

  pip install "pytorch-metric-learning[with-hooks]" -i https://pypi.tuna.tsinghua.edu.cn/simple
  

• CPU 版本：

  pip install "pytorch-metric-learning[with-hooks-cpu]" -i https://pypi.tuna.tsinghua.edu.cn/simple
  

Windows 用户特别注意： 如果在 Windows 上遇到 torch 兼容性问题，可能需要先指定版本安装 torch：

pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html
pip install pytorch-metric-learning

方式二：使用 Conda 安装

conda install -c conda-forge pytorch-metric-learning

注意：若需使用测试模块（Testers），需单独安装 faiss，可参考 FAISS 官方安装指南通过 conda 安装。

基本使用

该库的核心在于灵活组合 损失函数 (Losses) 和 挖掘器 (Miners)。以下是最基础的使用流程。

1. 基础用法：直接使用损失函数

最简单的场景是初始化一个损失函数（例如三元组损失 TripletMarginLoss），并在训练循环中传入模型生成的嵌入向量（embeddings）和标签（labels）。

from pytorch_metric_learning import losses

# 初始化损失函数
loss_func = losses.TripletMarginLoss()

# 训练循环示例
for i, (data, labels) in enumerate(dataloader):
    optimizer.zero_grad()
    
    # 获取嵌入向量 (形状: N x embedding_size)
    embeddings = model(data)
    
    # 计算损失 (labels 形状: N)
    loss = loss_func(embeddings, labels)
    
    loss.backward()
    optimizer.step()

原理：TripletMarginLoss 会根据标签自动在批次内构建所有可能的三元组（Anchor-Positive-Negative）。

2. 进阶用法：结合挖掘器 (Miner)

为了提升训练效果，通常只关注那些“难以区分”的样本对。可以使用 Miner 来筛选困难样本，并将其传递给损失函数。

from pytorch_metric_learning import miners, losses

# 初始化挖掘器和损失函数
miner = miners.MultiSimilarityMiner()
loss_func = losses.TripletMarginLoss()

# 训练循环示例
for i, (data, labels) in enumerate(dataloader):
    optimizer.zero_grad()
    embeddings = model(data)
    
    # 挖掘困难样本对
    hard_pairs = miner(embeddings, labels)
    
    # 将挖掘出的样本对传入损失函数
    # 库会自动将 pairs 转换为 triplets 进行计算
    loss = loss_func(embeddings, labels, hard_pairs)
    
    loss.backward()
    optimizer.step()

3. 自定义损失函数配置

您可以根据需求自定义距离计算方式、损失缩减策略以及正则化项。

from pytorch_metric_learning.distances import CosineSimilarity
from pytorch_metric_learning.reducers import ThresholdReducer
from pytorch_metric_learning.regularizers import LpRegularizer
from pytorch_metric_learning import losses

loss_func = losses.TripletMarginLoss(
    distance=CosineSimilarity(),        # 使用余弦相似度代替欧氏距离
    reducer=ThresholdReducer(high=0.3), # 丢弃大于 0.3 的损失值
    embedding_regularizer=LpRegularizer() # 添加 L2 正则化
)

4. 自监督学习支持

库提供了 SelfSupervisedLoss 包装器，用于自监督学习任务（如对比学习）。

from pytorch_metric_learning.losses import SelfSupervisedLoss

# 包装基础损失函数
loss_func = SelfSupervisedLoss(TripletMarginLoss())

for i, data in enumerate(dataloader):
    optimizer.zero_grad()
    
    # 原始数据嵌入
    embeddings = your_model(data)
    # 增强数据嵌入
    augmented = your_model(your_augmentation(data))
    
    # 计算自监督损失
    loss = loss_func(embeddings, augmented)
    
    loss.backward()
    optimizer.step()

更多高级功能（如内置训练器 Trainers、测试器 Testers 及精度计算器 AccuracyCalculator）请参考官方文档或 Google Colab 示例笔记。