0%

faiss 三种基础索引

前言

不同的索引方式是 faiss 的核心内容, 他们以不同的方式构建, 基于不同的算法与数据结构. 选择合适的 index 来处理数据是使用 faiss 最基础的一步. 官方 wiki 上也有帮助你如何选择不同的 index, 参见 Guidelines to choose an index
这次来学习 faiss 三个最基础的 index. 分别是 IndexFlatL2, IndexIVFFlat, IndexIVFPQ

IndexFlatL2 - 最基础的 Index

  1. 生成测试数据
    Faiss 处理固定维数 d 的向量集合,向量维度 d 通常为几十到几百。这些集合可以存储在矩阵中。我们假设行主存储,即。向量编号 i 的第 j 个分量存储在矩阵的第 i 行第 j 列中。Faiss 只使用 32 位浮点矩阵。

我们需要两个矩阵:

  • xb 对于数据库,它包含所有必须编入索引的向量,并且我们将搜索它。它的大小是 nb-by-d
  • xq 对于查询向量,我们需要找到最近的邻居。它的大小是 nq-by-d。如果我们只有一个查询向量,则 nq = 1。
    在下面的例子中,我们将使用在 d = 64 维度中形成均匀分布的矢量。为了测试也为了有趣, 我们在索引所依赖的第一维上做了一点小改变.
1
2
3
4
5
6
7
8
9
10
import numpy as np

d = 3
nb = 100000
nq = 10000
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.
  1. 构建索引并添加向量

Faiss 围绕 Index 物体构建。它们继承了一组向量库,并可以选择对它们进行预处理以提高搜索效率。Faiss 有很多类型的索引,我们将使用最简单的版本:IndexFlatL2,它只是对向量执行强力的 L2 距离搜索 (暴力搜索, brute-force)。

当索引被构建时, 他们都需要知道其所对应数据集向量的维度,也就是在我们的示例代码中的d. 然后,大多数索引还需要训练来分析向量集的分布。但是对于 IndexFlatL2 来说, 我们可以跳过这个操作。

当建立和训练完索引时,可以对索引执行两个操作:addsearch

将元素添加到索引。我们还可以输出索引的两个状态变量:

  • is_trained 表示索引是否需要训练的布尔值,
  • ntotal 索引中向量的数量。

一部分索引也可以存储对应于每个向量的整形 ID(但 IndexFlatL2 不行)。如果没有提供 ID,add 则使用向量序号作为 id,即。第一个向量是 0,第二个是 1 这样。

1
2
3
4
5
import faiss
index = faiss.IndexFlatL2(d)
print(index.is_trained)
index.add(xb)
print(index.ntotal)
  1. 在索引中搜索
    可以在索引上执行的基本搜索操作是 “k 临近搜索” 搜索,即。对于每个查询矢量,k 在向量集中查找其和他最近的向量。
    为了测试正确性,我们可以首先搜索几个已经存储的向量,来看一下它最近的邻居是否是这个向量本身。
1
2
3
4
5
6
7
k = 4
D, I = index.search(xb[:5], k)
print(I)
print(D)
D, I = index.search(xq, k)
print(I[:5])
print(I[-5:])
  1. 结果
    如果运行正常, 输出应该如图所示
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
True
100000
[[ 0 30 110 20]
[ 1 689 422 328]
[ 2 290 179 242]
[ 3 767 25 9]
[ 4 428 136 71]]
[[0.0000000e+00 1.1153221e-02 1.3206482e-02 1.4572144e-02]
[2.3841858e-07 2.8479099e-03 7.1740150e-03 9.6693039e-03]
[2.3841858e-07 7.0106983e-03 1.1301994e-02 1.7502308e-02]
[0.0000000e+00 2.8505325e-03 6.8078041e-03 1.2776375e-02]
[0.0000000e+00 3.4937859e-03 6.5484047e-03 8.7761879e-03]]
[[209 230 16 49]
[473 219 291 27]
[227 241 452 425]
[351 307 161 287]
[255 635 331 60]]
[[ 9782 10012 9661 9455]
[10236 10708 9918 9854]
[10716 10304 10381 10287]
[ 9414 9294 9890 9530]
[10306 10073 9710 10009]]
  • Ture 和 100000 表示索引是否经过训练及索引中的向量数
  • 中间一组结果是测试的结果 可以看到, 距离此向量最近的确实是他本身
  • 后面一组结果是搜索的结果

值得一提的是, IndexFlatL2 索引的结果是精确的, 可以用来作为其他索引测试中准确性程度的参考.

IndexFlatL2 完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


import numpy as np

d = 3
nb = 100000
nq = 10000
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

import faiss
index = faiss.IndexFlatL2(d)
print(index.is_trained)
index.add(xb)
print(index.ntotal)

k = 4
D, I = index.search(xb[:5], k)
print(I)
print(D)
D, I = index.search(xq, k)
print(I[:5])
print(I[-5:])

最后官方说在 2016 年的机器上这个 Demo 跑了大约 3.3 秒, 我自己测了一下跑了 1s 不到.

更快的搜索 - IndexIVFFlat

  1. 这太慢了,我怎么能让它更快?
    为了加快搜索速度,可以将数据集分割成几部分。我们在 d 维空间中定义 Voronoi 单元格,并且每个数据库矢量都落入其中一个单元格中。在搜索时,只有查询 x 所在单元中包含的数据库向量 y 与少数几个相邻查询向量进行比较。(划分搜索空间)

这是通过 IndexIVFFlat 索引完成的。这种类型的索引需要一个训练的过程,可以在与数据库向量具有相同分布的任何向量集合上执行。在这种情况下,我们只测试数据进行搜索。

这 IndexIVFFlat 还需要另一个索引,即量化器 (quantizer),它将矢量分配给 Voronoi 单元。每个单元由一个质心定义,找到一个矢量所在的 Voronoi 单元包括在质心集中找到该矢量的最近邻居。这是另一个索引的任务,通常是索引 IndexFlatL2。

搜索方法有两个参数:

  • nlist 划分单元的数量
  • nprobe 执行搜索访问的单元格数 (不包括 nlist)
  1. 结果
    nprobe = 1 时,结果看起来像
1
2
3
4
5
[[ 9900 10500  9831 10808]
[11055 10812 11321 10260]
[11353 10164 10719 11013]
[10571 10203 10793 10952]
[ 9582 10304 9622 9229]]

结果和上面的 L2 强力搜索类似,但是不同。这是因为一些结果可能不在完全相同的 Voronoi 单元格。因此,访问更多的单元格可能是有用的 (提高精度)。

nprobe 增加到 10 的确如此:

1
2
3
4
5
[[ 9900 10500  9309  9831]
[11055 10895 10812 11321]
[11353 11103 10164 9787]
[10571 10664 10632 9638]
[ 9628 9554 10036 9582]]

这是精确的结果。请注意,在这种情况下获得完美结果仅仅是数据分布的人为因素,因为它在 x 轴上具有强大的组件,这使得它更易于处理。nprobe 参数始终是调整结果速度和准确度之间折中的一种方式 。设置 nprobe = nlist 将给出与蛮力搜索(但会更慢)相同的结果。

IndexIVFFlat Demo 完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# encoding:utf-8

# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD+Patents license found in the
# LICENSE file in the root directory of this source tree.

# author : Facebook
# translate : h-j-13

import numpy as np
d = 64 # 向量维度
nb = 100000 # 向量集大小
nq = 10000 # 查询次数
np.random.seed(1234) # 随机种子,使结果可复现
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

import faiss

nlist = 100
k = 4
quantizer = faiss.IndexFlatL2(d) # the other index
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
# here we specify METRIC_L2, by default it performs inner-product search

assert not index.is_trained
index.train(xb)
assert index.is_trained

index.add(xb) # 添加索引可能会有一点慢
D, I = index.search(xq, k) # 搜索
print(I[-5:]) # 最初五次查询的结果
index.nprobe = 10 # 默认 nprobe 是1 ,可以设置的大一些试试
D, I = index.search(xq, k)
print(I[-5:]) # 最后五次查询的结果

更低的内存占用 - IndexIVFPQ

  1. 有损存储
    我们看到的索引IndexFlatL2IndexIVFFlat都存储完整的向量。 为了扩展到非常大的数据集,Faiss 提供了基于产品量化器的有损压缩来压缩存储的向量的变体。压缩的方法基于乘积量化 (Product Quantizer)
    在这种情况下,由于矢量没有精确存储,搜索方法返回的距离也是近似值

IndexIVFFlat 完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


import numpy as np

d = 64
nb = 100000
nq = 10000
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

import faiss

nlist = 100
m = 8
k = 4
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)

index.train(xb)
index.add(xb)
D, I = index.search(xb[:5], k)
print(I)
print(D)
index.nprobe = 10
D, I = index.search(xq, k)
print(I[-5:])
  1. 结果
1
2
3
4
5
6
7
8
9
10
11
[[   0  608  220  228]
[ 1 1063 277 617]
[ 2 46 114 304]
[ 3 791 527 316]
[ 4 159 288 393]]

[[ 1.40704751 6.19361687 6.34912491 6.35771513]
[ 1.49901485 5.66632462 5.94188499 6.29570007]
[ 1.63260388 6.04126883 6.18447495 6.26815748]
[ 1.5356375 6.33165455 6.64519501 6.86594009]
[ 1.46203303 6.5022912 6.62621975 6.63154221]]

我们可以观察到最近的邻居被正确地找到(它是矢量 ID 本身),但是向量自身的估计距离不是 0,尽管它远远低于与其他邻居的距离。这是由于有损压缩。

1
2
3
4
5
[[ 9432  9649  9900 10287]
[10229 10403 9829 9740]
[10847 10824 9787 10089]
[11268 10935 10260 10571]
[ 9582 10304 9616 9850]]

另外搜索真实查询时,虽然结果大多是错误的 (与刚才的 IVFFlat 进行比较),但是它们在正确的空间区域,而对于真实数据,情况更好,因为:

  • 统一数据很难进行索引,因为没有规律性可以被利用来聚集或降低维度
  • 对于自然数据,语义最近邻居往往比不相关的结果更接近。
  1. 简化指标结构
    由于构建索引可能会变得复杂,因此有一个工厂函数用于接受一个字符串来构造响应的索引。上面的索引可以通过以下简写获得:
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
2
3
4
5
6
7
num_list = 64
dim = 64
bytes_per_vector = 8
bits_per_sub_vector = 8

quantizer = faiss.IndexFlatL2(dim)
index = faiss.IndexIVFPQ(quantizer, dim, num_list, bytes_per_vector, bits_per_sub_vector)

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
2
3
4
5
index = faiss.IndexFlatL2(xb.shape[1])
ids = np.arange(xb.shape[0])
index.add_with_ids(xb, ids) # this will crash, because IndexFlatL2 does not support add_with_ids
index2 = faiss.IndexIDMap(index)
index2.add_with_ids(xb, ids) # works, the vectors are stored in the underlying index

参考文献

faiss wiki - Getting started
facebook Faiss 的基本使用示例 (逐步深入)