项目介绍:TensorFlow Similarity
项目背景
TensorFlow Similarity 是一个基于 TensorFlow 的库,专注于相似性学习。其技术包括自我监督学习、度量学习和对比学习。虽然这个库仍处在开发的 beta 阶段,因此可能会有一些重大变化,但它已经为研究、训练、评估和服务相似性模型及对比学习模型提供了最先进的算法和必要的组件。
项目概述
TensorFlow Similarity 提供了多种用于度量学习的工具,帮助用户研究各种数据集的相似性。它能帮助训练两大类模型:
-
自我监督模型:这种模型适用于无标签数据,通过学习这类数据的通用表示,从而提升少量标签的下游任务的准确率。用户可以先在大量无标签图片上预训练模型,然后在小规模有标签数据集上进行微调,从而达到更高的准确性。
-
相似性模型:这种模型输出的嵌入可以用来查找和聚类相似的例子,比如在一大批图片中找到代表同一对象的图片。通过训练相似性模型,用户可以找到如上所示的相似猫和狗的图片,即使只训练了一部分数据集的类别。
最新功能
- 在 2023 年 3 月发布了 0.17 版本,增加了更多的损失函数和度量,并进行了大规模重构:
- 添加了对比损失中的 VicReg 损失。
- 添加了用于检索论文中的度量,如 Precision@K。
- 本地支持分布式训练,比如 SimClr 现在能正确地在分布式训练环境下工作。
- 多模态嵌入的初始支持(CLIP)。
快速入门
安装
使用 pip 安装该库:
pip install --upgrade-strategy=only-if-needed tensorflow_similarity[tensorflow]
文档与示例
详细的示例笔记本是开始使用 TensorFlow Similarity 的好方法,这些笔记本可以在 Google Colab 上立即运行。
案例演示:MNIST 相似性模型
通过一个简单的示例展示如何在 MNIST 数据集上训练一个监督相似性模型。
数据准备
TensorFlow Similarity 提供了 数据采样器,它可以平衡批次,从而使得训练更加顺畅。
相似性模型构建
构建相似性模型类似于标准的 Keras 模型,但输出层通常是 MetricEmbedding()
层,以实现 L2 归一化,专门用于支持附加功能。
通过对比学习训练模型
使用相似性损失训练模型,输出可通过近似近邻搜索进行搜索的度量嵌入。
建立图片索引并查询
模型训练完成后,需要通过索引 API 将参考示例索引化,以进行搜索。
支持的算法
自我监督模型
- SimCLR
- SimSiam
- Barlow Twins
监督损失
- Triplet Loss
- PN Loss
- Multi Sim Loss
- Circle Loss
- Soft Nearest Neighbor Loss
度量
Tensorflow Similarity 提供常用的分类和检索评估度量,如 Precision、Recall 和 NDCG 等。
声明
这不是 Google 的官方产品。
如需在研究中使用 TensorFlow Similarity,请引用相关论文。
这一项目为用户提供了强大的功能和极大的灵活性,帮助他们在不同的数据集上进行度量学习和相似性分析。