ANN搜索—Faiss库的基本使用

ANN搜索—Faiss库的基本使用

Faiss 是一个用于高效相似性搜索和密集向量聚类的库。它包含搜索任意大小的向量集的算法，直到可能不适合 RAM 的向量集。它还包含用于评估和参数调整的支持代码。Faiss 是用 C++ 编写的，具有完整的 Python/numpy 包装器。一些最有用的算法是在 GPU 上实现的。

1. 安装

使用python调用Faiss的接口，官方推荐使用 conda 安装；

# CPU-only version
$ conda install -c pytorch faiss-cpu

# GPU(+CPU) version
$ conda install -c pytorch faiss-gpu

# or for a specific CUDA version
$ conda install -c pytorch faiss-gpu cudatoolkit=10.2 # for CUDA 10.2

2. 快速上手

Faiss处理固定维度的向量集合，这些向量集合一般以矩阵的形式存储，Faiss仅使用32位浮点数。

下面展示一个简单的示例

2.1 生成一些数据

生成一些维度为64的服从均匀分布向量，分别作为查询向量集和数据库向量集，然后使用查询向量在数据库向量中进行查询搜索。

import numpy as np

# In Python, the matrices are always represented as numpy arrays.
# The data type dtype must be float32.
dim = 64                        # 向量维度
database_size = 300000          # 数据库大小             
query_size = 10000              # 查询数量
np.random.seed(1234)            
vec_database = np.random.random((database_size, dim)).astype('float32')
vec_queries = np.random.random((query_size, dim)).astype('float32')

2.2 构建索引

Faiss通过 Index 对象来构建索引，它封装了数据库向量，并且可以对数据库向量进行预处理以提高搜索效率。Faiss 提供很多类型的 Index ，下面以最简单的版本 IndexFlatL2 为例，它表示对查询和数据库执行L2距离的暴力搜索。

示例：

import faiss

index = faiss.IndexFlatL2(dim)     # 传入向量维度以实例化一个 Index 对象

# Index对象具有两个属性，is_trained 和 ntotal
# is_trained 表示是否训练了的的bool状态值
# ntotal 表示索引的数量
index.add(vec_database)                  # 将数据库向量添加到 index 中
print(index.ntotal)                      # database_size
print(index.is_trained)                  # True

2.3 搜索

构建好索引后，就可以对查询向量进行近似 k 近邻搜索了，即在数据库向量集中查找到与查询向量的最近的k个向量。

搜索结果返回数据向量库中k个最近邻的向量索引 和 对应的k个距离分数（递增），以 numpy.ndarray 的数据类型返回。

k = 4                          # we want to see 4 nearest neighbors

distances, result_indexs = index.search(vec_queries, k)     # actual search
print(result_indexs[:5])                   # neighbors of the 5 first queries
print(result_indexs[-5:])                  # neighbors of the 5 last queries

print(type(distances))

3. 加快搜索速度

在 快速上手 中使用的是基于L2距离的暴力搜索，其实质就是计算查询向量和数据库中所有向量的L2距离，然后根据距离分数从小到大排序给出结果，并不是采用ANN算法进行搜索。Faiss 提供 IndexIVFFlat 对象来实现 ANN 搜索算法，关于ANN算法可参考下面链接：

https://towardsdatascience.com/comprehensive-guide-to-approximate-nearest-neighbors-algorithms-8b94f057d6b6

https://zhuanlan.zhihu.com/p/454511736

在 快速上手 示例中的查询向量和数据库向量的基础上采用ANN搜索，示例如下：

nlist = 100     # 类似于聚类, 将数据库中的向量聚类为100类
k = 4           # k 近邻

# IndexIVFFlat 需要传入一个 Index对象,作为ANN搜索算法采用的量化器 
quantizer = faiss.IndexFlatL2(dim)
# 创建一个应用ANN算法的 index 对象
index = faiss.IndexIVFFlat(quantizer, dim, nlist)

# 因为是采用ANN算法进行近似搜索,所以需要根据数据集中的向量集进行预处理,即将数据库所有向量聚类为 nlist 簇向量集
# 因此 index 的 is_trained 状态为 False
print(index.is_trained)
assert not index.is_trained

# 聚类预处理
index.train(vec_database)

print(index.is_trained)     # True
assert index.is_trained

# 将数据库中向量添加到 index 对象中
index.add(vec_database)

# 进行 ANN 搜索
distances, result_indexs = index.search(vec_queries, k)
print(result_indexs[:5])
print(result_indexs[-5:])

# IndexIVFFlat 的 nprobe 属性默认为1，在nprobe个最近邻的簇向量空间中进行 k 近邻搜索；
# 修改 nprobe 属性值，查看搜索结果
index.nprobe = 10
distances, result_indexs = index.search(vec_queries, k)
print(result_indexs[:5])
print(result_indexs[-5:])

index.nprobe = 20
distances, result_indexs = index.search(vec_queries, k)
print(result_indexs[:5])
print(result_indexs[-5:])

使用基于L2距离的暴力查询结果

nlist = 100, nprobe = 1 的 ANN搜索结果和暴力搜索结果对比

nlist = 100, nprobe = 10 的 ANN搜索结果和暴力搜索结果对比

nlist = 100, nprobe = 20 的 ANN搜索结果和暴力搜索结果对比

参考链接：

https://github.com/facebookresearch/faiss/wiki/