代码收藏家技术教程 2025-02-16

Python 向量检索库Faiss使用

Faiss（Facebook AI Similarity Search）是一个由 Facebook AI Research 开发的库，它专门用于高效地搜索和聚类大量向量。Faiss 能够在几毫秒内搜索数亿个向量，这使得它非常适合于实现近似最近邻（ANN）搜索，这在许多应用中都非常有用，比如图像检索、推荐系统和自然语言处理。

以下是如何使用 Faiss 的基本步骤和示例：

1. 安装 Faiss

首先，你需要安装 Faiss。你可以使用 pip 来安装它：

pip install faiss-cpu  # 对于 CPU 版本
# 或者
pip install faiss-gpu  # 对于 GPU 版本（需要 CUDA 支持）

注意：安装 GPU 版本时，请确保你的系统已经安装了 CUDA，并且 CUDA 的版本与 Faiss 兼容。

2. 导入 Faiss

在你的 Python 脚本中导入 Faiss：

import faiss

3. 准备数据

你需要准备一些向量数据来进行检索。这些向量可以是任何维度的，但通常它们是从特征提取器中得到的，比如深度学习模型的输出。

# 假设我们有一些 128 维的向量
nb, d = 10000, 128  # 10000 个向量，每个向量 128 维
xb = np.random.random((nb, d)).astype('float32')

4. 构建索引

使用 Faiss 提供的索引类型之一来构建索引。有多种索引类型可供选择，每种类型在构建时间、搜索速度和准确性方面都有所不同。

# 创建一个 HNSW 索引（适用于 GPU 和 CPU）
index = faiss.IndexHNSWFlat(d, 16)  # d 是向量维度，16 是 M 参数（影响性能和准确性）
# 或者使用其他索引类型，如 IndexFlatL2（适用于 CPU）
# index = faiss.IndexFlatL2(d)

# 将向量添加到索引中
index.add(xb)

5. 搜索向量

现在你可以使用索引来搜索与给定查询向量最接近的向量了。

# 准备一些查询向量
nq = 5  # 查询向量的数量
xq = np.random.random((nq, d)).astype('float32')

# 搜索与每个查询向量最接近的 k 个向量
k = 4  # 返回最接近的 k 个向量
D, I = index.search(xq, k)  # D 是距离数组，I 是索引数组

# 打印结果
print("查询向量与最近邻的距离:")
print(D)
print("最近邻的索引:")
print(I)