faiss 三个最基础的 index. 分别是 IndexFlatL2, IndexIVFFlat, IndexIVFPQ
搜索时,可以以查询向量为中心,返回距离在一定范围内的结果,如返回数据库中与查询向量距离小于0.3的结果。不是所有的Index都支持按距离检索,但是下面三种Index都支持,只支持在CPU使用。
一、IndexFlatL2 - 最基础的Index
IndexFlatL2索引的结果是精确的,可以用来作为其他索引测试中准确性程度的参考.
import numpy as np import faiss import time def demo_IndexFlatL2(): d = 2048 # 2048维的数据 nb = 2 # database size nq = 1 # nb of queries np.random.seed(1234) index = faiss.IndexFlatL2(d) print(index.is_trained) # 随机nb个数据插入索引 xb = np.random.random((nb, d)).astype('float32') xb[:, 0] += np.arange(nb) / 1000. index.add(xb) # 再随机nb个数据插入索引 xb1 = np.random.random((nb, d)).astype('float32') xb1[:, 0] += np.arange(nb) / 1000. index.add(xb1) print('index.ntotal:', index.ntotal) # 随机nq个数据用于查询 xq = np.random.random((nq, d)).astype('float32') xq[:, 0] += np.arange(nq) / 1000. k = 4 # 查询距离最近的4个 D, I = index.search(xq, k) # 返回值D是距离 I是索引 print("I: ", I) # 按照距离查询 dist = 1000 # 定义一个半径/阈值 _, D, I = index.range_search(xb[[2], :], dist) # 用第2个向量查询 print('range_search res:', I) # 删除元素,dtype=np.int64非常重要 print('remove之前ntotal:', index.ntotal) index.remove_ids(np.asarray((2,3), dtype=np.int64)) print('remove之后ntotal:', index.ntotal) if __name__ == '__main__': demo_IndexFlatL2()
二、更快的搜索 - IndexIVFFlat
为了加快搜索速度,可以将数据集分割成几部分。我们在d维空间中定义Voronoi单元格,并且每个数据库矢量都落入其中一个单元格中。在搜索时,只有查询x所在单元中包含的数据库向量y与少数几个相邻查询向量进行比较。(划分搜索空间)
这种类型的索引需要一个训练的过程,可以在与数据库向量具有相同分布的任何向量集合上执行。
这IndexIVFFlat还需要另一个索引,即量化器(quantizer),它将矢量分配给Voronoi单元。每个单元由一个质心定义,找到一个矢量所在的Voronoi单元包括在质心集中找到该矢量的最近邻居。这是另一个索引的任务,通常是索引IndexFlatL2。
搜索方法有两个参数:
- nlist 划分单元的数量
- nprobe 执行搜索访问的单元格数(不包括nlist)
nprobe
参数始终是调整结果速度和准确度之间折中的一种方式 。设置 nprobe = nlist 将给出与蛮力搜索(但会更慢)相同的结果。
import numpy as np import faiss import time def demo_IndexIVFFlat(): d = 64 # 向量维度 nb = 200 # 向量集大小 nq = 10000 # 查询次数 np.random.seed(1234) # 随机种子,使结果可复现 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 xb = np.random.random((nb, d)).astype('float32') xb[:, 0] += np.arange(nb) / 1000. index.train(xb) index.add(xb) # 添加索引可能会有一点慢 print('index.ntotal:', index.ntotal) xb1 = np.random.random((nb, d)).astype('float32') xb1[:, 0] += np.arange(nb) / 1000. index.train(xb1) index.add(xb1) # 添加索引可能会有一点慢 print('index.ntotal:', index.ntotal) xq = np.random.random((nq, d)).astype('float32') xq[:, 0] += np.arange(nq) / 1000. D, I = index.search(xq, k) # 搜索 print(I[-5:]) # 最初五次查询的结果 index.nprobe = 10 # 执行搜索访问的单元格数 默认 nprobe 是1 ,可以设置的大一些试试 D, I = index.search(xq, k) print(I[-5:]) # 最后五次查询的结果 # 按照距离查询 dist = 1000 # 定义一个半径/阈值 _, D, I = index.range_search(xb[[2], :], dist) # 用第2个向量查询 print('range_search res:', I) # 删除元素,dtype=np.int64非常重要 print('remove之前ntotal:', index.ntotal) index.remove_ids(np.asarray((2,3), dtype=np.int64)) print('remove之后ntotal:', index.ntotal) if __name__ == '__main__': demo_IndexIVFFlat()
三、更低的内存占用 - IndexIVFPQ
索引IndexFlatL2
和IndexIVFFlat
都存储完整的向量。 为了扩展到非常大的数据集,Faiss提供了基于产品量化器的有损压缩来压缩存储的向量的变体。压缩的方法基于乘积量化(Product Quantizer)。
在这种情况下,由于矢量没有精确存储,搜索方法返回的距离也是近似值。
import numpy as np import faiss import time def demo_IndexIVFPQ(): 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. nlist = 100 m = 8 k = 4 quantizer = faiss.IndexFlatL2(d) # 内部的索引方式依然不变 index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) # 每个向量都被编码为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:]) # 按照距离查询 dist = 1000 # 定义一个半径/阈值 _, D, I = index.range_search(xb[[2], :], dist) # 用第2个向量查询 print('range_search res:', I) # 删除元素,dtype=np.int64非常重要 print('remove之前ntotal:', index.ntotal) index.remove_ids(np.asarray((2,3), dtype=np.int64)) print('remove之后ntotal:', index.ntotal) if __name__ == '__main__': demo_IndexIVFPQ()