faiss 三种基础索引方式

        faiss 三个最基础的 index. 分别是 IndexFlatL2IndexIVFFlatIndexIVFPQ

        搜索时,可以以查询向量为中心,返回距离在一定范围内的结果,如返回数据库中与查询向量距离小于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

        索引IndexFlatL2IndexIVFFlat都存储完整的向量。 为了扩展到非常大的数据集,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()

猜你喜欢

转载自blog.csdn.net/xian0710830114/article/details/121767389