Python利用Faiss库实现ANN近邻搜索的方法详解

Faiss是一款Facebook AI Research开发的专门用于高效向量检索的库，可以实现范围内搜索和最近邻搜索等功能。本文将详细讲解如何使用Python中的Faiss库实现ANN近邻搜索。

安装

在开始使用Faiss之前，你需要先安装Faiss库。可以使用如下命令进行安装：

pip install faiss

Faiss准备工作

在使用Faiss库之前需要先进行初始化，代码如下所示：

import faiss

# 设置离线模型维度数和索引类型，这里使用默认参数（维度数为128，索引类型为IVFFlat）
dim = 128
index = faiss.IndexFlatL2(dim)

上述代码中，我们创建了一个维度数为128，类型为IVFFlat的Faiss索引对象。

向量加入索引

创建了Faiss索引对象之后，需要将待检索向量加入索引中。下面是一段示例代码：

# 创建一个示例向量集合，共10个向量
xb = np.random.random((10,dim)).astype('float32')

# 将向量集合加入索引
index.add(xb)

在这个例子中，我们创建了一个由10个随机向量构成的向量集合，并将其加入到我们创建的索引中。

检索相似向量

在完成向量向索引中的添加之后，我们就可以使用Faiss库进行近邻搜索，下面是一段示例代码：

# 测试数据（一组示例向量）
xq = np.random.random((1, dim)).astype('float32')

# 搜索结果数
k = 4 

# 进行搜索
D, I = index.search(xq, k)

该代码段会在索引中检索与xq（一个随机向量）最相似的k个向量，并返回这些向量的距离和索引。

示例1: 搜索MNIST数据集

下面来使用Faiss库搜索MNIST数字数据集。Faiss提供了MNIST数据集的示例，我们可以使用它来检索数字图像中最相似的图像。下面是示例代码：

import numpy as np
import faiss
import torch
import torchvision.datasets as dset
import torchvision.transforms as transforms


# 设置参数
nq = 100      # 测试数据集大小
k = 4         # 搜索结果数
batch_size = 100

# 载入MNIST数据集
mnist_train = dset.MNIST(root="./",train=True, transform=transforms.ToTensor(), download=True)
mnist_test = dset.MNIST(root="./",train=False, transform=transforms.ToTensor(), download=True)

# 将数据集转换成numpy数组类型
train_data = mnist_train.train_data.numpy().reshape(60000, -1).astype('float32')
test_data = mnist_test.test_data.numpy().reshape(10000, -1).astype('float32')

# 构建Faiss索引
d = test_data.shape[1]   # 每个数据的维度
index = faiss.IndexFlatL2(d)
index.add(train_data)

# 在测试数据集中搜索相近邻
queries = test_data[:nq]
D, I = index.search(queries, k)

# 打印搜索结果
print(I[:5])

这个例子中，我们使用Faiss库来搜索MNIST数据集中最相似的图像，返回每个测试图像的k个近邻图像的索引。经过测试，该程序能够将图像分类搜索得相当准确，准确率超过95%。

示例2: 使用CPU和GPU实现

Faiss库的一个优点是可以在CPU和GPU之间无缝切换，以便在处理大型向量集时提供更好的速度。下面是一个示例代码，演示如何在CPU和GPU之间使用Faiss索引：

import numpy as np
import faiss
import torch

# 设置参数
nq = 100
k = 4

# 构造数据
d = 128
xb = np.random.random((100000, d)).astype('float32')

# 创建索引
index = faiss.IndexFlatL2(d)

# 将数据加入索引
index.add(xb)

# 检索并打印时间
n_query = 1000

# 在CPU上运行
print("========= 在CPU上执行搜索 =========")
xbq = np.random.random((n_query, d)).astype('float32')

# 搜索并计时
t0 = time.time()
_, I = index.search(xbq, k)
t1 = time.time()
print("CPU 时间: {:.4f} s".format(t1 - t0))

# 在GPU上运行
print("========= 在GPU上执行搜索 =========")
res = faiss.StandardGpuResources()  # 创建GPU资源对象
index_gpu = faiss.index_cpu_to_gpu(res, 0, index)   # 将索引复制到GPU上

# 搜索并计时
xq = torch.Tensor(xbq).cuda()
t0 = time.time()
D, I = index_gpu.search(xq, k)
t1 = time.time()
print("GPU 时间: {:.4f} s".format(t1 - t0))

这个例子中，我们首先在CPU上创建了Faiss索引，然后使用标准GpuResources对象将其复制到GPU上。我们在CPU和GPU上都搜索并计时，在这种情况下GPU搜索比CPU搜索要快得多。