前言
不同的索引方式是 faiss 的核心内容, 他们以不同的方式构建, 基于不同的算法与数据结构. 选择合适的 index 来处理数据是使用 faiss 最基础的一步. 官方 wiki 上也有帮助你如何选择不同的 index, 参见 Guidelines to choose an index
这次来学习 faiss 三个最基础的 index. 分别是 IndexFlatL2
, IndexIVFFlat
, IndexIVFPQ
IndexFlatL2 - 最基础的 Index
- 生成测试数据
Faiss 处理固定维数 d 的向量集合,向量维度 d 通常为几十到几百。这些集合可以存储在矩阵中。我们假设行主存储,即。向量编号 i 的第 j 个分量存储在矩阵的第 i 行第 j 列中。Faiss 只使用 32 位浮点矩阵。
我们需要两个矩阵:
- xb 对于数据库,它包含所有必须编入索引的向量,并且我们将搜索它。它的大小是 nb-by-d
- xq 对于查询向量,我们需要找到最近的邻居。它的大小是 nq-by-d。如果我们只有一个查询向量,则 nq = 1。
在下面的例子中,我们将使用在 d = 64 维度中形成均匀分布的矢量。为了测试也为了有趣, 我们在索引所依赖的第一维上做了一点小改变.
1 | import numpy as np |
- 构建索引并添加向量
Faiss 围绕 Index 物体构建。它们继承了一组向量库,并可以选择对它们进行预处理以提高搜索效率。Faiss 有很多类型的索引,我们将使用最简单的版本:IndexFlatL2,它只是对向量执行强力的 L2 距离搜索 (暴力搜索, brute-force)。
当索引被构建时, 他们都需要知道其所对应数据集向量的维度,也就是在我们的示例代码中的d
. 然后,大多数索引还需要训练来分析向量集的分布。但是对于 IndexFlatL2 来说, 我们可以跳过这个操作。
当建立和训练完索引时,可以对索引执行两个操作:add 和 search。
将元素添加到索引。我们还可以输出索引的两个状态变量:
is_trained
表示索引是否需要训练的布尔值,ntotal
索引中向量的数量。
一部分索引也可以存储对应于每个向量的整形 ID(但 IndexFlatL2 不行)。如果没有提供 ID,add 则使用向量序号作为 id,即。第一个向量是 0,第二个是 1 这样。
1 | import faiss |
- 在索引中搜索
可以在索引上执行的基本搜索操作是 “k 临近搜索” 搜索,即。对于每个查询矢量,k 在向量集中查找其和他最近的向量。
为了测试正确性,我们可以首先搜索几个已经存储的向量,来看一下它最近的邻居是否是这个向量本身。
1 | k = 4 |
- 结果
如果运行正常, 输出应该如图所示
1 | True |
- Ture 和 100000 表示索引是否经过训练及索引中的向量数
- 中间一组结果是测试的结果 可以看到, 距离此向量最近的确实是他本身
- 后面一组结果是搜索的结果
值得一提的是, IndexFlatL2 索引的结果是精确的, 可以用来作为其他索引测试中准确性程度的参考.
IndexFlatL2 完整代码
1 |
|
最后官方说在 2016 年的机器上这个 Demo 跑了大约 3.3 秒, 我自己测了一下跑了 1s 不到.
更快的搜索 - IndexIVFFlat
- 这太慢了,我怎么能让它更快?
为了加快搜索速度,可以将数据集分割成几部分。我们在 d 维空间中定义 Voronoi 单元格,并且每个数据库矢量都落入其中一个单元格中。在搜索时,只有查询 x 所在单元中包含的数据库向量 y 与少数几个相邻查询向量进行比较。(划分搜索空间)
这是通过 IndexIVFFlat 索引完成的。这种类型的索引需要一个训练的过程,可以在与数据库向量具有相同分布的任何向量集合上执行。在这种情况下,我们只测试数据进行搜索。
这 IndexIVFFlat 还需要另一个索引,即量化器 (quantizer),它将矢量分配给 Voronoi 单元。每个单元由一个质心定义,找到一个矢量所在的 Voronoi 单元包括在质心集中找到该矢量的最近邻居。这是另一个索引的任务,通常是索引 IndexFlatL2。
搜索方法有两个参数:
- nlist 划分单元的数量
- nprobe 执行搜索访问的单元格数 (不包括 nlist)
- 结果
当nprobe
= 1 时,结果看起来像
1 | [[ 9900 10500 9831 10808] |
结果和上面的 L2 强力搜索类似,但是不同。这是因为一些结果可能不在完全相同的 Voronoi 单元格。因此,访问更多的单元格可能是有用的 (提高精度)。
将 nprobe
增加到 10 的确如此:
1 | [[ 9900 10500 9309 9831] |
这是精确的结果。请注意,在这种情况下获得完美结果仅仅是数据分布的人为因素,因为它在 x 轴上具有强大的组件,这使得它更易于处理。nprobe
参数始终是调整结果速度和准确度之间折中的一种方式 。设置 nprobe = nlist 将给出与蛮力搜索(但会更慢)相同的结果。
IndexIVFFlat Demo 完整代码
1 | # encoding:utf-8 |
更低的内存占用 - IndexIVFPQ
- 有损存储
我们看到的索引IndexFlatL2
和IndexIVFFlat
都存储完整的向量。 为了扩展到非常大的数据集,Faiss 提供了基于产品量化器的有损压缩来压缩存储的向量的变体。压缩的方法基于乘积量化 (Product Quantizer)。
在这种情况下,由于矢量没有精确存储,搜索方法返回的距离也是近似值。
IndexIVFFlat 完整代码
1 |
|
- 结果
1 | [[ 0 608 220 228] |
我们可以观察到最近的邻居被正确地找到(它是矢量 ID 本身),但是向量自身的估计距离不是 0,尽管它远远低于与其他邻居的距离。这是由于有损压缩。
1 | [[ 9432 9649 9900 10287] |
另外搜索真实查询时,虽然结果大多是错误的 (与刚才的 IVFFlat 进行比较),但是它们在正确的空间区域,而对于真实数据,情况更好,因为:
- 统一数据很难进行索引,因为没有规律性可以被利用来聚集或降低维度
- 对于自然数据,语义最近邻居往往比不相关的结果更接近。
- 简化指标结构
由于构建索引可能会变得复杂,因此有一个工厂函数用于接受一个字符串来构造响应的索引。上面的索引可以通过以下简写获得:
1 | index = faiss.index_factory(d,“ IVF100,PQ8 ”) |
更换 PQ4 用 Flat 得到的 IndexFlat。当预处理(PCA)应用于输入向量时,工厂特别有用。例如,预处理的工厂字符串通过 PCA 投影将矢量减少到 32 维为:“PCA32,IVF100,Flat”。
总结
1,支持两种相似性计算方法:L2 距离(即欧式距离)和点乘(归一化的向量点乘即 cosine 相似度);
2,按照是否编码压缩数据可以分为两类算法,使用压缩的算法可以在单台机器上处理十亿级别的向量规模;
3,并非线程安全的——不支持并行添加向量或搜索与添加的并行;仅在 CPU 模式下支持并行搜索;
4,只有继承了 IndexIVF 的算法才支持向量的 remove() 操作,但由于是连续存储,remove 的时间复杂度是 O(n),建议另外维护一个列表记录被删除的或尚存的向量;
5,faiss 针对批量搜索做了优化;
6,IndexPQ, IndexIVFFlat, IndexIVFPQ, IndexIVFPQR 需要训练;
7,不支持重新训练,建议新建一个索引;
8,只接受 32-bit 浮点类型的输入数据;
9,使用 index = faiss.index_factory(dim, “PCA32,IVF100,PQ8”) 这种形式创建索引更灵活,此处类型参数可解释为:使用 PCA 降维将原始向量降至 32 维,使用 IVF 建立索引,子 list(即 bucket 分桶)个数为 100,使用 Product Quantizer (乘积量化) 将每个向量压缩编码成 8 字节;等价于
1 | num_list = 64 |
10,索引类型的选择
- 如果需要精确的搜索结果,不要降维、不要量化,使用 Flat,同时,使用 Flat 意味着数据不会被压缩,将占用同等大小的内存;
- 如果内存很紧张,可以使用 PCA 降维、PQ 量化编码,来减少内存占用,最终占用的内存大小约等于 <降维后的向量维度> <量化后的每个向量的字节数> <向量个数>; 如果量化编码后的字节数大于 64,推荐使用 SQx 替换 PQx,准确度相同但速度会更快;为了便于量化编码,可以使用 OPQx_y 先对向量做线性变换,y 必须是编码后字节数 x 的倍数,但最好小于维度 dim 和 4*x;
- 如果总向量个数 N 小于 1 百万,推荐使用 IVFx ,x 的选值介于 4sqrt(N) 和 16sqrt(N) 之间,训练数据的大小至少要是 x 的 30 倍;如果总向量个数 N 大于 1 百万、小于 1 千万,推荐使用 IMI2x10,实际内部聚类个数是 2 ^ (2 * 10),将需要 64 * 2 ^ 10 个向量参与训练;如果总向量个数 N 大于 1 千万、小于 1 亿,推荐使用 IMI2x12;如果总向量个数 N 大于 1 亿、小于 10 亿,推荐使用 IMI2x14;IMI 方法不支持 GPU;
- IndexIVF 天生支持 add_with_ids 方法,对于不支持 add_with_ids 方法的类型,可以使用 IndexIDMap 辅助
1 | index = faiss.IndexFlatL2(xb.shape[1]) |