向量数据库Faiss的搭建与使用(python)

发布于:2024-08-26 ⋅ 阅读:(137) ⋅ 点赞:(0)

1. 安装 Faiss
1.1 通过 pip 安装
如果你使用的是 Python 环境,可以通过 pip 安装 Faiss 的 Python 绑定:

pip install faiss-cpu

如果你需要 GPU 支持,还需要安装 CUDA 和 cuDNN,并安装相应的 Faiss 版本:

pip install faiss-gpu

1.2 从源代码编译
如果你需要定制化编译,可以从 GitHub 获取源代码并自行编译:

git clone https://github.com/facebookresearch/faiss.git
cd faiss
make
make python
sudo make install
sudo make python_install

2. 使用 Faiss
2.1 准备数据
假设你已经有一些向量数据,这些数据可以是图像特征向量、文本嵌入向量等。这里我们创建一些随机向量作为示例:

import numpy as np
from faiss import IndexFlatL2

# 创建一个 10000 x 128 的随机向量矩阵
d = 128  # 向量维度
nb = 10000  # 向量数量
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.

2.2 构建索引
接下来,我们将使用 Faiss 构建一个 L2 距离的索引,并将向量数据添加到索引中:

index = IndexFlatL2(d)   # 创建一个 L2 距离的索引
print(index.is_trained)
index.add(xb)                  # 添加向量数据
print(index.ntotal)

2.3 搜索向量
现在我们可以使用 Faiss 来搜索最近邻向量:

# 创建一些查询向量
nq = 1000
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

# 搜索最近的 k 个向量
k = 4
D, I = index.search(xq, k)  # 实际上返回的是距离和索引
print(I[:5])               # 输出最近的 5 个索引
print(D[:5])               # 输出最近的 5 个距离

3. 更高级的索引类型
Faiss 支持多种索引类型,例如 IVF (Inverted File) 索引,可以进一步提高搜索效率:

from faiss import IndexIVFFlat

nlist = 100
quantizer = IndexFlatL2(d)  # 量化器
index_ivf = IndexIVFFlat(quantizer, d, nlist, IndexFlatL2.METRIC_L2)
index_ivf.train(xb)
index_ivf.add(xb)

D, I = index_ivf.search(xq, k)
print(I[:5])
print(D[:5])

4. 使用 GPU
如果你的系统支持 GPU,可以使用 GPU 版本的 Faiss 来加速搜索:

from faiss import StandardGpuResources, GpuIndexFlatL2

res = StandardGpuResources()  # 使用 GPU 资源
gpu_index = GpuIndexFlatL2(res, d)

gpu_index.add(xb)
D, I = gpu_index.search(xq, k)
print(I[:5])
print(D[:5])