TensorFlow Similarity: 人类的度量学习
TensorFlow Similarity 是一个 TensorFlow 库,用于 相似性学习,包括自监督学习、度量学习、相似性学习和对比学习等技术。TensorFlow Similarity 仍处于测试阶段,我们可能会进行破坏性更改。
介绍
TensorFlow Similarity 提供了最先进的度量学习算法以及进行研究、训练、评估和服务基于相似性和对比模型所需的所有组件。这些组件包括模型、损失函数、度量、采样器、可视化工具和索引子系统,使这一过程快速且简便。
使用 TensorFlow Similarity 您可以训练两种主要类型的模型:
-
自监督模型:用于在未标记数据上学习通用数据表示,从而提高下游任务的准确性。例如,您可以使用 TensorFlow Similarity 支持的某种对比方法在大量未标记图像上预训练模型,然后在少量标记数据集上进行微调以达到更高的准确性。要开始训练自己的自监督模型,请参见此 笔记本。
-
相似性模型:输出嵌入,允许您在大量示例中查找和聚类类似的示例,例如表示相同对象的图像。例如,如上所见,您可以训练一个相似性模型,在只训练少量数据集类的情况下,从 Oxford IIIT Pet Dataset 中查找和聚类相似的,看起来未见过的猫和狗图片。要开始训练自己的相似性模型,请参见此 笔记本。
更新内容
- [2023年3月]: 0.17 版增加了更多损失和度量,并进行了大规模重构
- 对比损失中增加了 VicReg 损失。
- 添加了检索论文中使用的度量,例如 Precision@K。
- 原生支持分布式训练,例如 SimClr 现在支持分布式训练。
- 初步支持多模态嵌入 (CLIP)。
有关更多详细信息和先前版本信息,请参阅 更新日志
入门指南
安装
使用 pip 安装该库。
注意:如果你已经安装了 tensorflow>=2.4,可以省略 TensorFlow extra_require 关键字。
pip install --upgrade-strategy=only-if-needed tensorflow_similarity[tensorflow]
文档
详细且有解说的 笔记本 是快速上手 TensorFlow Similarity 的好方法。这里可能有一个与您的数据或问题类似的例子(如果没有,请告诉我们)。您可以点击 Google Colab 图标立即在 Google Colab 上使用这些示例。
有关具体函数的更多信息,您可以 查看 API 文档
如果要为该项目贡献,请查看 贡献指南
最简示例:MNIST 相似性
点击展开并查看如何使用 TF.Similarity 在 mnist 数据上训练一个监督相似性模型
下面是一个简单的示例,演示了如何在 MNIST 数据集上训练 TensorFlow Similarity 模型。此示例演示了 TensorFlow Similarity 提供的一些主要组件及其如何协同工作。请参考 hello_world 笔记本 以了解更详细的介绍。
准备数据
TensorFlow Similarity 提供了 数据采样器,适用于各种数据集类型,平衡批次以确保训练更加顺利。 在此示例中,我们使用了直接从 TensorFlow 数据集目录中集成的多拍采样器。
from tensorflow_similarity.samplers import TFDatasetMultiShotMemorySampler
# 从 MNIST 数据集生成平衡批次的数据采样器
sampler = TFDatasetMultiShotMemorySampler(dataset_name='mnist', classes_per_batch=10)
构建相似性模型
构建 TensorFlow Similarity 模型类似于构建标准 Keras 模型,不同的是输出层通常是一个 MetricEmbedding()
层,它执行 L2 归一化,且模型被实例化为支持附加功能的特殊子类 SimilarityModel()
。
from tensorflow.keras import layers
from tensorflow_similarity.layers import MetricEmbedding
from tensorflow_similarity.models import SimilarityModel
# 使用标准 Keras 层构建相似性模型
inputs = layers.Input(shape=(28, 28, 1))
x = layers.experimental.preprocessing.Rescaling(1/255)(inputs)
x = layers.Conv2D(64, 3, activation='relu')(x)
x = layers.Flatten()(x)
x = layers.Dense(64, activation='relu')(x)
outputs = MetricEmbedding(64)(x)
# 构建一个特殊的相似性模型
model = SimilarityModel(inputs, outputs)
通过对比学习训练模型
要输出一个可通过近似最近邻搜索进行检索的度量嵌入,模型需要使用相似性损失进行训练。在这里我们使用 MultiSimilarityLoss()
,这是最有效的损失函数之一。
from tensorflow_similarity.losses import MultiSimilarityLoss
# 使用对比损失训练相似性模型
model.compile('adam', loss=MultiSimilarityLoss())
model.fit(sampler, epochs=5)
构建图像索引并查询
一旦模型经过训练,必须通过模型索引 API 索引参考示例以进行检索。索引后,您可以使用模型查找 API 搜索索引中K个最相似的项。
from tensorflow_similarity.visualization import viz_neigbors_imgs
# 索引100个嵌入的 MNIST 示例以便进行检索
sx, sy = sampler.get_slice(0,100)
model.index(x=sx, y=sy, data=sx)
# 为给定示例查找最相似的五个索引 MNIST 示例
qx, qy = sampler.get_slice(3713, 1)
nns = model.single_lookup(qx[0])
# 可视化查询示例及其前五个邻居
viz_neigbors_imgs(qx[0], qy[0], nns)
支持的算法
自监督模型
- SimCLR
- SimSiam
- Barlow Twins
有监督的损失函数
- Triplet Loss
- PN Loss
- Multi Sim Loss
- Circle Loss
- Soft Nearest Neighbor Loss
度量
TensorFlow Similarity 提供了许多用于 分类 和 检索 评估的常用度量。包括:
名称 | 类型 | 描述 |
---|---|---|
Precision | 分类 | |
Recall | 分类 | |
F1 Score | 分类 | |
Recall@K | 检索 | |
Binary NDCG | 检索 |
引用
如果在您的研究中使用 TensorFlow Similarity 的任何部分,请引用此参考文献:
@article{EBSIM21,
title={TensorFlow Similarity: A Usable, High-Performance Metric Learning Library},
author={Elie Bursztein, James Long, Shun Lin, Owen Vallis, Francois Chollet},
journal={Fixme},
year={2021}
}
免责声明
这不是谷歌的官方产品。