PyTorch中的深度度量学习

学习用于图像检索或其他信息检索任务的深度度量学习方法。

我们的XBM算法被提名CVPR 2020最佳论文。

知乎上关于XBM的一篇文章

我在知乎上写了一篇文章，通俗易懂地快速解读了XBM的思想和动机：

跨越时空的难样本挖掘

欢迎大家阅读并提出宝贵意见！

推荐一篇最近发表的、非我撰写的优秀DML论文：

度量学习的现实检验

由康奈尔科技与Facebook AI联合发布

摘要：过去四年来，深度度量学习领域的论文不断宣称在准确率方面取得了巨大进展，常常使性能比十年前的方法提升一倍以上。本文对这一领域进行了深入分析，以验证这些说法是否属实。我们发现这些论文在实验设计上存在缺陷，并提出了一种新的评估度量学习算法的方法。最终，我们的实验结果表明，随着时间的推移，性能的提升充其量也只是微不足道的。

　　

XBM：一种新的SOTA DML方法，已被CVPR-2020接受为口头报告，并被提名为最佳论文：

用于嵌入学习的跨批次记忆机制

• 显著提升：XBM可在三个大规模数据集上将R@1指标提升12%至25%

• 易于实现：仅需几行代码即可完成

• 内存高效：处理大规模数据集时内存占用不到1GB

• 代码已公开：xbm

其他实现：

pytorch-metric-learning(由Kevin Musgrave开发的优秀项目)

　　 　　

基于GPW的MS损失：已被CVPR 2019接受为海报展示

具有通用成对权重的多相似性损失用于深度度量学习

• 代码公开链接

• 论文新版本 ，为了让大家更容易理解我的想法，我最近重写了论文的主要部分，使其更加清晰。（2020年3月24日）

　

本仓库中实现的深度度量学习方法：

• 对比损失 [1]

• 半难样本挖掘策略 [2]

• 提升结构损失* [3] （因其原始性能较弱而进行了改进）

• 二项式BinDeviance损失 [4]

• NCA损失 [6]

• 多相似性损失 [7]

数据集

• Car-196

  前98类作为训练集，后98类作为测试集

• CUB-200-2011

  前100类作为训练集，后100类作为测试集

• 斯坦福在线产品

  用于实验的部分，我们将11,318个类别的59,551张图片划分为训练集，11,316个类别的60,502张图片划分为测试集。

• In-Shop服装检索

  对于In-Shop服装检索数据集，包含3,997个类别的25,882张图片用于训练。 测试集则分为查询集（3,985个类别的14,218张图片）和画廊集（3,985个类别的12,612张图片）。

• 预处理过的CUB和Cars196

  提取码：inmj

  为了便于复现实验效果，我提供了预处理过的CUB和Cars-196数据集。

环境要求

• Python >= 3.5
• PyTorch = 1.0

在CUB-200和Cars-196上的SOTA对比

Recall@K	1	2	4	8	16	32	1	2	4	8	16	32
HDC	53.6	65.7	77.0	85.6	91.5	95.5	73.7	83.2	89.5	93.8	96.7	98.4
聚类	48.2	61.4	71.8	81.9	-	-	58.1	70.6	80.3	87.8	-	-
ProxyNCA	49.2	61.9	67.9	72.4	-	-	73.2	82.4	86.4	87.8	-	-
智能挖掘	49.8	62.3	74.1	83.3	-	-	64.7	76.2	84.2	90.2	-	-
Margin [5]	63.6	74.4	83.1	90.0	94.2	-	79.6	86.5	91.9	95.1	97.3	-
HTL	57.1	68.8	78.7	86.5	92.5	95.5	81.4	88.0	92.7	95.7	97.4	99.0
ABIER	57.5	68.7	78.3	86.2	91.9	95.5	82.0	89.0	93.2	96.1	97.8	98.7

在SOP和In-shop上的SOTA对比

Recall@K	1	10	100	1000	1	10	20	30	40	50
聚类	67.0	83.7	93.2	-	-	-	-	-	-	-
HDC	69.5	84.4	92.8	97.7	62.1	84.9	89.0	91.2	92.3	93.1
Margin [5]	72.7	86.2	93.8	98.0	-	-	-	-	-	-
Proxy-NCA	73.7	-	-	-	-	-	-	-	-	-
ABIER	74.2	86.9	94.0	97.8	83.1	95.1	96.9	97.5	97.8	98.0
HTL	74.8	88.3	94.8	98.4	80.9	94.3	95.8	97.2	97.4	97.8

更多细节请参阅我们在CVPR-2019发表的论文多相似性损失

复现Car-196（或CUB-200-2011）实验

*** 权重 :***

sh run_train_00.sh

其他实现：

[TensorFlow] (由geonm提供)

参考文献

[1] [R. Hadsell, S. Chopra, and Y. LeCun. 通过学习不变映射进行降维]

[2] [F. Schroff, D. Kalenichenko, and J. Philbin. FaceNet：用于人脸识别和聚类的统一嵌入。CVPR，2015年。]

[3][H. Oh Song, Y. Xiang, S. Jegelka, and S. Savarese. 通过提升结构化特征嵌入进行深度度量学习。CVPR，2016年。]

[4][D. Yi, Z. Lei, and S. Z. Li. 面向实际应用的人体再识别的深度度量学习。]

[5][C. Wu, R. Manmatha, A. J. Smola, and P. Kr¨ahenb¨uhl. 深度嵌入学习中的采样问题。ICCV，2017年。]

[6][R. Salakhutdinov and G. Hinton. 通过保持类别邻域结构来学习非线性嵌入。AISTATS，2007年。]

引用方式

如果您在研究中使用了本方法或代码，请按以下方式引用：

@inproceedings{wang2019multi,
title={Multi-Similarity Loss with General Pair Weighting for Deep Metric Learning},
author={Wang, Xun and Han, Xintong and Huang, Weilin and Dong, Dengke and Scott, Matthew R},
booktitle={CVPR},
year={2019}
}

@inproceedings{wang2020xbm,
title={Cross-Batch Memory for Embedding Learning},
author={Wang, Xun and Zhang, haozhi and Huang, Weilin and Scott, Matthew R},
booktitle={CVPR},
year={2020}
}

Deep_Metric 快速上手指南

Deep_Metric 是一个基于 PyTorch 的深度学习度量学习（Deep Metric Learning, DML）开源库，主要用于图像检索等任务。该仓库实现了多种主流损失函数（如 Multi-Similarity Loss），并包含了在 CVPR 2020 获最佳论文提名的 XBM (Cross-Batch Memory) 方法。

环境准备

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

• 操作系统: Linux / macOS / Windows
• Python: >= 3.5
• 深度学习框架: PyTorch == 1.0 (或兼容版本)
• 其他依赖: 建议安装 torchvision 以加载标准数据集

提示：国内用户可使用清华源或阿里源加速 Python 包的安装。

安装步骤

1. 克隆项目代码

将仓库克隆到本地：

git clone https://github.com/MalongTech/research-ms-loss.git
cd research-ms-loss

2. 安装依赖

虽然 README 未提供 requirements.txt，但核心依赖仅为 PyTorch。您可以使用以下命令安装（推荐国内镜像源）：

pip install torch==1.0.0 torchvision==0.2.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 准备数据集

项目支持 Car-196, CUB-200-2011, Stanford-Online-Products (SOP), In-Shop Clothes 等数据集。

• 自动下载：部分脚本可能支持自动下载，需检查具体数据集加载代码。
• 手动下载（推荐）：作者提供了处理好的 CUB 和 Cars-196 数据集以便复现结果。
  • 下载地址：百度网盘链接
  • 提取码：inmj
  • 解压后请将数据放置在项目指定的 data 目录下（具体路径请参考代码中的 dataset 配置）。

基本使用

本仓库最核心的功能是训练深度度量学习模型。以下以复现 Car-196 或 CUB-200-2011 数据集上的实验为例。

运行训练脚本

项目根目录提供了封装好的 Shell 脚本，可直接启动训练流程。这将默认使用 Multi-Similarity Loss (MS Loss) 或其他配置的损失函数进行训练。

sh run_train_00.sh

代码结构简析

• 损失函数实现: 在 losses 相关模块中，包含 Contrastive Loss, Triplet Loss (Semi-Hard), Lifted Structure Loss, Binomial Deviance Loss, NCA Loss 以及核心的 Multi-Similarity Loss。
• XBM 模块: 跨批次记忆模块代码位于 ret_benchmark/modeling/xbm.py，可轻松集成到现有模型中以提升 Recall@1 指标（提升幅度约 12%~25%）。
• 模型评估: 训练完成后，脚本会自动在测试集上计算 Recall@K 指标（K=1, 2, 4, 8, 16, 32 等）。

自定义开发

若需在自有数据集上使用 XBM 或 MS Loss，可参考 ret_benchmark 目录下的建模逻辑，仅需几行代码即可将 CrossBatchMemory 嵌入到您的 PyTorch 训练循环中。