Faiss (Facebook AI Similarity Search) 是一个用于高效相似向量搜索的库，特别适合处理大规模向量数据。

  • 它的核心功能是给定一个向量数据库，为当前的query vector(s) 寻找Top-K个相似向量。
  • 此外也支持KNN聚类，PCA降维，数据量化等模块。
  • 官方文档：https://github.com/facebookresearch/faiss/wiki
1
2
3
4
5

# 安装方式: 一个环境不能同时安装cpu与gpu两个版本
pip install faiss-cpu

pip install faiss-gpu
# 下面以cpu版本为例

1. 相似Top-K查找

Faiss通过构建一个Index（索引）对象，用于建立可实现快速相似搜索的数据结构。

常见的索引类型：

  • IndexFlatL2: 暴力搜索，精度高，速度慢（L2距离）
  • IndexFlatIP: 暴力搜索，适用于内积（dot product）
  • IndexIVFFlat: IVF (Inverted File Indexing) 可加速搜索
  • IndexIVFPQ: PQ (Product Quantizer) 可降低缓存需求
  • IndexHNSWFlat: HNSW (Hierarchical Navigable Small World): fast and accurate

Basic indexes 介绍了全部的index类型；

Guidelines to choose an index 介绍下选择index的参考标准。

1.1 暴力搜索

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

import numpy as np
import faiss

# (1) 初始化index
d = 64          # 需要提前设置向量的维度
index = faiss.IndexFlatL2(d)
print(index.is_trained)
# True
# 对于暴力搜索方式，不需要训练，因此这里为True。但对于其它方式需要手动进行训练


# (2) 添加向量数据库
nb = 100000
xb = np.random.random((nb, d)).astype('float32')
xb.shape
# (100000, 64)
index.add(xb)   
print(index.ntotal)
# 100000


# (3) 对query向量查询数据库的Top-K相似向量
k = 3
D, I = index.search(xq, k) 
I.shape  
# (10000, 4): Top-K的索引
I[:3]
# array([[74531, 40931, 30745, 68436],
#        [53407, 38558, 47694, 51554],
#        [82772, 20841,  3102, 25325]])

D.shape    
# (10000, 4): Top-K的度量值
D[:3]
# array([[4.651325 , 5.663311 , 5.687874 , 5.708172 ],
#        [4.663849 , 5.218609 , 5.2473526, 5.313629 ],
#        [5.563755 , 5.568184 , 5.588524 , 5.8358154]], dtype=float32)

1.2 IVF方式举例

IVF通过事先将数据库的向量分为若干个类群（Voronoi cells），然后计算query vector与每个类群质心的距离度量，可降低搜索范围，加快速度。其中有两个关键参数

  • nlist: 分为多少个类群
  • nprobe: 最终选取多少Top类群，作为搜索范围
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import numpy as np
import faiss

# (1) 初始化index
nlist = 100
k = 4
quantizer = faiss.IndexFlatL2(d)  
index = faiss.IndexIVFFlat(quantizer, d, nlist)
index.train(xb) #训练


# (2) 添加向量数据库
index.add(xb)  
print(index.ntotal)
# 100000


# (3) 对query向量查询数据库的Top-K相似向量
k = 3
D, I = index.search(xq, k) 
I[:3]
# array([[42962, 51466, 57950],
#        [51817, 59530, 89886],
#        [95115, 79896, 39241]])

index.nprobe = 10              # default nprobe is 1, try a few more
D, I = index.search(xq, k)
I[:3]
# array([[40931, 30745, 96040],
#        [53407, 38558, 47694],
#        [82772, 25325, 27231]])

暴力搜索是最精确的方式，其它所有方式都会或多或少损失一定精度，以最求其它方面的提升（速度/显存）。PS：每种方式都有自己的参数。

1.3 联合datasets使用

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

from datasets import Dataset
import pandas as pd
import numpy as np

## (1) 示例数据集
texts = ["s1", "s2", "s3", "s4", "s5"]
embeddings = np.random.rand(5, 64).tolist()

df = pd.DataFrame({"text": texts, "embeddings": embeddings})
dataset = Dataset.from_pandas(df)
# Dataset({
#     features: ['text', 'embeddings'],
#     num_rows: 5
# })

## (2) 创建索引
dataset.add_faiss_index(column="embeddings")
# 等价于
dataset.add_faiss_index(column="embeddings", index_factory="Flat")
# 可以设置其它值(e.g. "IVF100,PQ64")，具体可参加文档

## (3) 对query向量进行查询
query_embedding = np.random.rand(64,)
scores, retrieved_examples = dataset.get_nearest_examples(
    "embeddings",  # embedding 列名
    query_embedding,  # query embedding
    k=3  # top k
)
retrieved_examples["text"]
# ['s3', 's5', 's1']


## (4) 保存与加载

# 保存 Dataset 和 faiss 索引
dataset.save_faiss_index("embeddings", "my_faiss.index")
dataset.drop_index("embeddings")
dataset.save_to_disk("my_dataset")

# 加载 Dataset 和索引
from datasets import load_from_disk
dataset = load_from_disk("my_dataset")
dataset.load_faiss_index("embeddings", "my_faiss.index")

2. 聚类降维

2.1 K-means聚类

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import faiss
import numpy as np

x = np.random.rand(1000, 128).astype('float32')
ncentroids = 16

niter = 20
verbose = True
d = x.shape[1]
kmeans = faiss.Kmeans(d, ncentroids, niter=niter, verbose=verbose)
kmeans.train(x)

# Clustering 1000 points in 128D to 16 clusters, redo 1 times, 20 iterations
#   Preprocessing in 0.00 s
#   Iteration 19 (0.48 s, search 0.15 s): objective=10058 imbalance=1.136 nsplit=0         
# 10057.9833984375

# 所有类群的质心Embedding
kmeans.centroids.shape
# (16, 128)

# 计算每个vector属于哪个类群
D, I = kmeans.index.search(x, 1)

I.shape
# (1000, 1)
I[:4]
# array([[ 2],
#        [11],
#        [14],
#        [ 5]])

2.2 PCA降维

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# random training data 
mt = np.random.rand(1000, 40).astype('float32')
mat = faiss.PCAMatrix (40, 10)
mat.train(mt)
assert mat.is_trained
tr = mat.apply(mt)
tr.shape
# (1000, 10)

# print this to show that the magnitude of tr's columns is decreasing
print (tr ** 2).sum(0)  # 每列的方差递减