faiss 三个最基础的 index. 分别是 IndexFlatL2
, IndexIVFFlat
, IndexIVFPQ
搜索时,可以以查询向量为中心,返回距离在一定范围内的结果,如返回数据库中与查询向量距离小于0.3的结果。不是所有的Index都支持按距离检索,但是下面三种Index都支持,只支持在CPU使用。
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()
为了加快搜索速度,可以将数据集分割成几部分。我们在d维空间中定义Voronoi单元格,并且每个数据库矢量都落入其中一个单元格中。在搜索时,只有查询x所在单元中包含的数据库向量y与少数几个相邻查询向量进行比较。(划分搜索空间)
这种类型的索引需要一个训练的过程,可以在与数据库向量具有相同分布的任何向量集合上执行。
这IndexIVFFlat还需要另一个索引,即量化器(quantizer),它将矢量分配给Voronoi单元。每个单元由一个质心定义,找到一个矢量所在的Voronoi单元包括在质心集中找到该矢量的最近邻居。这是另一个索引的任务,通常是索引IndexFlatL2。
搜索方法有两个参数:
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()
索引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()