Faiss(Facebook AI Similarity Search)向量数据库是由Facebook AI研究院开发的一种高效相似性搜索和聚类的库。Faiss不仅支持在高维空间中进行高效的相似性搜索,还能够在处理大规模数据集时展现出卓越的性能,尤其适用于图像检索、文本搜索、推荐系统和语音处理等多种应用场景。
- 在推荐系统中,Faiss可以用于快速查找用户可能感兴趣的物品或寻找具有相似兴趣的用户。
- 在信息检索领域,Faiss可以用于构建文档或图像的相似性搜索引擎。
- 在图像识别领域,Faiss可以用于构建图像特征的索引,实现快速的相似图像搜索和图像聚类。
安装Faiss:
💢cpu版本:
conda install -c pytorch faiss-cpu
💢gpu版本:
conda install -c pytorch faiss-gpu
Faiss 处理固定维数 d 的向量集合,通常为几十到几百个。这些集合可以存储在矩阵中。我们假设采用行主存储,即向量编号 i 的第 j 个分量存储在矩阵的第 i 行、第 j 列中。Faiss 仅使用 32 位浮点矩阵。
import numpy as np
d = 64 # 设置向量的维度为64
nb = 100000 # 向量数量为100,000
nq = 10000 # 查询向量的数量为10,000
np.random.seed(1234) # make reproducible
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. # 修改查询向量的第一个维度
建立索引
Faiss 是围绕Index
对象构建的。它封装了一组数据库向量,并可选地对它们进行预处理,以提高搜索效率。索引有很多种类型,我们将使用最简单的版本,即对它们进行强力的 L2 距离搜索:IndexFlatL2
。
所有索引在构建时都需要知道它们所操作的向量的维数,当索引建立并训练完成后,可以对索引进行两种操作:add
和search
。
- 当我们说一个索引是否被“训练”时,我们实际上是指该索引是否已经通过某种方式优化了其内部结构,以便更有效地处理搜索查询。
import faiss
index = faiss.IndexFlatL2(d) # 创建一个IndexFlatL2类型的索引
print(index.is_trained) # 打印出索引是否已经被训练
index.add(xb) # 将向量添加到索引中
print(index.ntotal)
IndexFlatL2
索引是一种简单的暴力搜索索引,它不需要训练过程,因为它直接计算查询向量与数据库中所有向量的L2距离,以找到最相似的向量。
搜索~
可以在索引上执行的基本搜索操作是k
最近邻搜索,即对于每个查询向量,k
在数据库中找到其最近的邻居。
此操作的结果可以方便地存储在大小为nq
-by-的整数矩阵中k
,其中第 i 行包含查询向量 i 的邻居的 ID,按距离递增排序。除了这个矩阵之外,该search
操作还返回一个nq
-by-k
浮点矩阵,其中包含相应的平方距离。
k = 4 # we want to see 4 nearest neighbors
D, I = index.search(xb[:5], k) # 在索引中搜索xb数组的前5个向量(xb[:5])的k个最近邻居
print(I)
print(D)
D, I = index.search(xq, k) # 整个查询集xq上搜索每个查询向量的k个最近邻居
print(I[:5]) # 前5个查询向量的最近邻居的索引位置
print(I[-5:])
D
:包含了查询向量与其找到的最近邻居之间的距离。I
:也是一个数组,但它包含的是最近邻居在索引中的位置或索引。
结果:
💥由于索引中未添加任何向量,因此无法进行有效的相似性搜索。在实际应用中,我们需要先将向量添加到索引中,然后才能进行搜索操作。
💥向索引添加向量:
nb = 100000 # 假设有100,000个向量
xb = np.random.random((nb, d)).astype('float32') # 生成随机向量数据,100000个64维数据
index.add(xb) # 将向量数据添加到索引中
# 优化索引(跳过)
💯结果:
- 后两个为实际的搜索输出(前五和后五)。
更快的搜索!
为了加快搜索速度,可以将数据集分割成块。我们在 d 维空间中定义 Voronoi 单元,每个数据库向量都位于其中一个单元中。在搜索时,仅将查询 x 所在的单元中包含的数据库向量 y 和一些相邻的向量与查询向量进行比较。
这是通过IndexIVFFlat
索引完成的。这种类型的索引需要一个训练阶段,可以对具有与数据库向量相同分布的任何向量集合执行。
还IndexIVFFlat
需要另一个索引,即量化器,它将向量分配给 Voronoi 单元。每个单元由一个质心定义,找到向量所在的 Voronoi 单元就是在质心集合中找到向量的最近邻居。这是另一个索引的任务,通常是IndexFlatL2
。
nlist = 100
# nlist指定IndexIVFFlat索引中聚类中心的数量
k = 4
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist)
# 这个索引中,quantizer 被用作内部机制来量化向量,并将它们分配到倒排文件中的不同聚类中心
# assert 语句用于验证索引的状态,确保其在训练前后的行为符合预期。
assert not index.is_trained
index.train(xb)
print("~~训练完成~~")
assert index.is_trained
index.add(xb)
D, I = index.search(xq, k)
print(I[-5:])
index.nprobe = 10 # 在搜索时控制要检查的聚类中心的数量
D, I = index.search(xq, k)
print(I[-5:])
# quantizer 被“嵌入”到 index 中,是因为 index 需要使用 quantizer 的量化功能来将向量分配到正确的聚类中心,并实现高效的搜索。
⭐️IndexIVFFlat
索引首先将向量空间划分为nlist
个聚类中心,然后使用quantizer
(IndexFlatL2
索引)来量化这些中心。- ⭐️在搜索时,
IndexIVFFlat
索引会先确定查询向量所属的聚类中心,然后只在该中心的向量中执行搜索,从而大大减少了计算量。 ⭐️xb
(代表数据集的一部分或全部)来训练索引。这是为了优化量化器或聚类中心。- ⭐️我们重新初始化了索引 index,所以第一次添加的操作就没有影响了。
- 🌟每创建一个索引,就相当于在向量搜索的上下文中创建了一个独立的、用于存储和查询向量的数据结构。
搜索方法有两个参数:nlist
,即单元格数量,以及nprobe
,即执行搜索时访问的单元格数量(共nlist
)
🧊nprobe = 10:
🧊nprobe = 1 :
💢设置 nprobe = nlist 会得到与强力搜索相同的结果(但速度较慢)。
Faiss支持将索引保存到磁盘文件中,并在需要时重新加载它们。通过保存和重新加载索引,可以在不同的会话或应用程序中重用索引:
# 保存索引
faiss.write_index(index, 'index.faiss')
# 加载索引
index = faiss.read_index('index.faiss')