faiss 三种基础索引方式

小灰灰 2023-10-01 08:40 241阅读 0赞

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()

[Product Quantizer]: https://hal.archives-ouvertes.fr/file/index/docid/514462/filename/paper_hal.pdf