Faiss 源码剖析：类结构分析

摘要：在下文中，我将尝试通过 Faiss 源码中各种类结构的设计来梳理 Faiss 中的各种概念以及它们之间的关系。

本文分享自华为云社区《Faiss源码剖析（一）:类结构分析》，原文作者：HW007。

Faiss 是由 Facebook AI Research 研发的为稠密向量提供高效相似度搜索和聚类的框架。通过其官方给出的新手指南，我们可以快速地体验 Faiss 的基本功能。但是，相信大多数人看完官方的新手指南后，对 Faiss 很多的概念还是有点模糊、无法清晰的明确这些概念之间的边界。比如说在 Faiss 中，Quantizer 是个什么概念、其与 Index 之间的联系是什么；还有各种 Index 之间的关系又是什么等等。为此，在下文中，我将尝试通过 Faiss 源码中各种类结构的设计来梳理 Faiss 中的各种概念以及它们之间的关系。

首先奉上 Faiss 源码的类图全家福如下，详细的 EA 类图文件见附件：

图一：Faiss 的类图全家福

首先，我们来看一下 Faiss 最主要的功能：相似度搜索。如下图所示，以图片搜索为例，所谓相似度搜索，便是在给定的一堆图片（下图中左上角的图集）中，寻找出我指定的目标（下图中左下角的巴士图片）最像的 K 张图片，也简称为 KNN（K 近邻）问题。

接下来我们看一下为了解决 KNN 问题，在工程上我们至少需要做哪些事情。显然，有两件事是必须要做的，第一，我们要把上面例子中的那个图库存储起来；第二，当用户指定一种图片后，我们需要知道怎么从存储的图库中找到最近相似的 K 张图片。由此，我们确定了 Faiss 在其应用场景中至少应该具备的两个功能：添加功能和搜索功能。

对于熟悉数据库的同学来说，应该能在这里嗅到点“CRUD”的味道。的确，当我们对“图集”有添加存储这样的动作后，修改和删除等功能也便接踵而来了。由此 Faiss 本质上就是一个向量数据库。对于数据库来说，时空优化是两个永恒的主题，即在存储上如何以更少的空间来存储更多的信息，在搜索上如何以更快的速度来搜索出更准确的信息。如何减少搜索所需的时间？在数据库中很最常见的操作便是加各种索引，把各种加速搜索算法的功能或空间换时间的策略都封装成各种各样的索引，以满足各种不同的引用场景。

由此，我们便不难理解为什么 Faiss 中为什么会有那么多的 Index 了，因为 Index 这个概念本身就与加速搜索是绑在一起的。由此也可以看出在 Faiss 中，如何又快又准地找到相似向量是第一要务。下图中给出的是 Faiss 中最重要的两个基类：Index 和 IndexBinary。

在上图中，用白色的箭头标出了这两个基类中最重要的三个函数，其中 add()和 search() 函数便对应了我上文中所提到的 Faiss 至少应该实现的两个基本功能：存储和搜索。在此顺带提一下，与传统的数据库相比，Faiss 的 Index 还包含了数据存储的功能，如果你一开始就从字面上按照传统数据库中索引的概念来理解地话，就会感觉有点怪怪的。接下来，我们重点聊聊 Index 中的 train()函数，我们都知道天上是不会白白掉馅饼的，对于 Faiss 来说，不管其为了减少存储空间还是加速搜索，都需要提前做好一些准备工作，这便是 train()函数发挥作用的时候了。

以减少存储为例子，我们都知道在图片处理中通过 PCA 可以将图片从高维空间（p 维）转换到低维空间（q 维, 其中 p > q ），其具体操作便是是将高维空间中的图片向量（n*p）乘以一个转换矩阵（p*q），得到一个低维空间中的向量（n*q）。为了使得在整个降维的过程中信息丢失最少，我们需要对待转换图片进行分析计算得到相应的转换矩阵（p*q）。也就是说这个降维中乘以的转换矩阵是与待转换图片息息相关的。

回到我们的 Faiss 中来，假设我期望使用 PCA 预处理来减少 Index 中的存储空间，那在整个处理流程中，除了输入搜索图库外，我必须多输入一个转换矩阵，但是这个转换矩阵是与图库息息相关的，是可以由图库数据计算出来的。如果把这个转换矩阵看成一个参数的话，我们可以发现，在 Faiss 的一些预处理中，我们会引入一些参数，这些参数又无法一开始由人工来指定，只能通过喂样本来训练出来，所以 Index 中需要有这样的一个 train() 函数来为这种参数的训练提供输入训练样本的接口。由此，我们也可以发现，这些喂给 train()函数的样本数据最好与之后要添加存储的图集以及搜索目标一致比较好，比如说，你先给 Index 喂一个猪脸数据集训练出 PCA 中的转换矩阵，再给这个 Index 添加人脸数据集，最后再在这个索引上做人脸识别，这样肯定比不上一开始就喂人脸数据集得到 PCA 转换矩阵的效果好。

由上，我们已经可以从 train()、add()和 search()三大函数大概地了解到 Faiss 中的 Index 是个什么东西了，接下来我们看一下 Faiss 中有哪些不同的 Index。从图一中的类图中可以看到，在 Faiss 中，大多数类基本都继承或使用了 Index 接口，他们要么对 Index 接口中定义的 train、add 和 search 函数进行了自己个性化的实现（如图一中被淡橙色标注的类），要么就是对已经实现的三大函数的类进行包装，提供一些三大函数之外的流程上的加工处理（如图一中被淡蓝色标注的类）。

从图一中我们可以看到这些被淡蓝色标注的偏包装的 Index 子类，他们与 Index 基类之间既有“is a”又有“hold a”关系，在类结构上出现这种关系的时候，设计者要么是在设计一个树或链表的节点，要么是在设计一个包装类。显然在 Faiss 中更偏向于后者。一方面，淡蓝色的 Index 子类借助其所“hold”的 Index 来提供基本的 train、add 和 search 功能，使其自身符合 Index 接口的定义标准，成为一种 Index，为之后的层层嵌套包装提供支持。另一方面，他又对其所“hold”的 Index 类进行了一些通用的功能扩展。如下图的 IndexPreTransform 类所示，Faiss 将对待存储图集的预处理，如归一化、PCA 降维等功能抽象成一个 VectorTransform 接口，让 IndexPreTransform 使用它来为其所“hold”的 Index 添加预处理功能，这种预处理功能是与其所“hold”的是什么 Index 没有任何关系，因此我更偏向于将这种功能归结为 Index 之外的流程上的包装功能。如 IndexPreTransform 类提供了数据预处理功能、IndexIDMap 类提供了自定义 ID 功能、IndexShards 类为 Index 的并行计算提供了相关的支持等。

接下来我们来看一下图一中被淡橙色标注的 Index 子类，如 IndexLSH、IndexPQ、IndexIVFPQ 等，从名字中我们可以大概了解到这些类都是基于一些不同的算法实现的不同索引，他们的 train、add 和 search 方法各有差异。但在整体上还是能找到一些其他结构上的共性。在上文中，我们知道 Index 具有存储的功能，这些被淡橙色标注的 Index 子类在数据存储方式上基本可以划分为两大类，一类是统一存到一个容器中，如在 IndexLSH、IndexPQ 等中我们都可以看到一个命名为 codes 的 vector 容器。另一类是分桶储存到多个容器中，这主要为索引后续的非精确分桶局部搜索提供支持，为此，Faiss 特地抽象出 InvertedLists 接口，需要支持分桶局部搜索的 Index 子类均会有 hold 一个实现了 InvertedLists 接口（淡紫色标注）的实例来存储其数据。如下图所示，Faiss 为 InvertedLists 接口提供了数组、链表和磁盘文件等三种不同的实现。

在图一中还有两个被标记为淡绿色的类 ProductQuantizer 和 ScalarQuantizer 值得大家关注下，在结构上，这两个类均没有派生的子类，并且所有其他的类与他们的关系均为“hold a”关系，很纯粹的工具类。从其命名中的 Quantizer（量化器）后缀可知，这两个工具类的作用是将“连续或稠密”的数据进行“离散或稀疏化”，简单来说就是进行聚类的操作，就像我们把 18 岁以下的称为少年，18~50 岁的称为中年一样，我们把具体年龄量化成年龄段的过程就是一个聚类的过程。从图一中还可以看到，带有 Quantizer 后缀的类还有四个：MultiIndexQuantizer、MultiIndexQuantizer2、IndexScalarQuantizer 和 Level1Quantizer。其中前三个均是通过对 ProductQuantizer 或 ScalarQuantizer 的包装来实现 Quantizer 的功能，没什么稀奇的地方，但最后一个 Level1Quantizer 类竟然是包装了两个 Index 类，而且其中一个 Index 类的属性名还是 quantizer，如下图所示。

难道 Index 也是一种 Quantizer？的确，对于 Index 来说，我们更熟悉的是其将数据集存储起来，再寻找某个数据在该数据集中的 K 个最近邻点的功能。但如果 Index 中存储的是数据分类后各个类的中心点呢，那么对于某个数据，我们便可以在该 Index 上通过 KNN 来求得其 K（此时 K=1）个最近邻点，这些求出来的中心点所代表的类便是该数据在聚类中该归属的类。由此我们可以看到 Index 是可用来聚类，将数据量化成类的中心点的。因此，Index 可以被包装成一个 Quantizer 也便不足为奇了。其实 Index 的这种聚类功能在 Faiss 的设计中是很常见的，除了上面所说的用来做 Quantizer 外，还可以用来辅助实现 K-means 算法，这也是为什么 Level1Quantizer 类中除 quantizer 外还存在一个名为 clustering_index 的 Index 类型属性的原因。通过上面的分析，我们还可以知道，在 Faiss 的 Quantizer 类中，或明或暗都应该有个地方来存储用来辅助量化的“centroids”，即类中心点，它们在大多数场景中都是经过数据训练出来的（如对数据进行 K-means 聚类），在少数场景中也可以直接人为设定。

让我们最后来关注下 IndexIVF 类（上图中被圈出来的淡紫色类）。也许在上文介绍淡紫色的 InvertedLists 类簇时，有人会有疑问，InvertedLists 类及其派生子类在 Faiss 中主要为 Index 提供非精确的分桶局部搜索功能，这种功能与 Index 的种类毫无关系，按上文对 Index 派生的子类的分类标准来看，IndexIVF 类应该是一个偏包装的 Index 子类，应该被标注为淡蓝色才对。的确，如上图所示，虽然 IndexIVF 类没有直接“hold a”Index 类，但其通过继承 Level1Quantizer 类间接“hold a”Index 类，确实也是一个偏包装的 Index 派生子类。图一的颜色标注只是为了突出拥有 IVF 功能的 Index 类，通过颜色来辅助各个功能类簇在视觉上的区分度而已，不必深究。

通过上文，我们可以发现，Faiss 的整个类结构设计是非常清晰简洁的，其首先将 KNN 问题的解决过程切分成 train、add 和 search 三个步骤并抽象出 Index 基类。接着从这些基类派生出各种偏功能实现或者偏流程包装的 Index 子类。此外还为 Index 提供了两种的存储方式：集中和分桶（IVF）。最后还提供了 SQ 和 PQ 两种量化编码工具以及将这些编码工具或其他的 Index 包装成 Quantizer 的类。

点击关注，第一时间了解华为云新鲜技术~