Faiss学习笔记
目录
- Faiss概念
- 使用步骤
- 常见索引类型及示例
-
- IndexFlatL2
- IndexIVFFlat
- IndexIVFPQ
- 索引选择指南
- 分布式检索
- 总结
- 参考文档
Faiss概念
Faiss是一个向量检索库。它的GitHub上描述是这样的
Faiss is a library for efficient similarity search and clustering of dense vectors. It contains algorithms that search in sets of vectors of any size, up to ones that possibly do not fit in RAM. It also contains supporting code for evaluation and parameter tuning. Faiss is written in C++ with complete wrappers for Python/numpy. Some of the most useful algorithms are implemented on the GPU. It is developed primarily at Meta’s Fundamental AI Research group.
通过以上描述，可以发现有如下关键点：
- 支持稠密向量相似度搜索和聚类
- 包含了可以搜索任意维度向量集合的算法
- 使用C++编写，支持Python/numpy接口
- 部分最有用的算法支持GPU
本质上是解决基于向量相似度的knn问题
使用步骤
简单来说就是构建索引、添加数据、查询数据。具体步骤为：
- 构建索引index
- 根据不同索引的特性，对索引进行训练（train）。对于有些算法来说，训练的本质是聚类；对于IndexFlatL2等可以直接搜索的场景可省略本步骤。
- add 添加完整数据集到索引
- 针对query数据进行搜索search操作
常见索引类型及示例
根据特点，索引类别主要分为几类。完整版参见文档：
- Flat indexes（暴力检索类，搜索时逐个比对）
- IndexIVF* indexes（IVF，Inverted File System基于分区的方法，例如聚类。此方法在检索时只会与一部分向量做对比，提升了效率。但可能得到的结果落入局部最优而非全局最优。在构建索引时还可以采用乘积量化、降维等方式进行内存压缩和提速）
- IndexHNSW及相关变体（HNSW，Hierarchical Navigable Small World，基于索引数据构建图，检索时根据生成图进行最邻近检索。HNSW算法介绍可以参见这篇）
- IndexLSH（LSH，局部敏感哈希）
以下为几个常见方法简介
# test_datasets.py
# 数据集
import numpy as np
import time
print("generate data ..")
a = time.time()
d = 64 # dimension
nb = 1000000 # database size
nq = 10000 # nb of queries
np.random.seed(1234) # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.
print("generate data done, time cost:{}".format(time.time() - a))
IndexFlatL2
基于L2距离的暴力全量搜索，速度较慢
这里多介绍一下L1距离和L2距离
L1距离：即曼哈顿距离。每一维度坐标相减后的绝对值求和。公式为
d 1 ( v 1 , v 2 ) = ∑ i = 1 n ∣ x i − y i ∣ d_1(v_1,v_2)= \sum_{i=1}^n|x_i-y_i| d1(v1,v2)=i=1∑n∣xi−yi∣
L2距离：即欧氏距离。
d 2 ( v 1 , v 2 ) = ∑ i = 1 n ( x i − y i ) 2 d_2(v_1,v_2)=\sqrt{\sum_{i=1}^{n}{\left( x_{i}-y_{i} \right)^{2}}} d2(v1,v2)=i=1∑n(xi−yi)2
代码如下
from test_dataset import *
import faiss
import numpy as np
import time
print("start to build index..")
a = time.time()
index = faiss.IndexFlatL2(d) # build the index
print("build index done, time cost:{}".format(time.time() - a ))
print("is trained: {}".format(index.is_trained)) # IndexFlatL2不需要训练过程
print("start to add query..")
a = time.time()
index.add(xb) # add vectors to the index
print("add data done, time cost:{}".format(time.time() - a ))
print(index.ntotal)
print("start to search..")
k = 4 # we want to see 4 nearest neighbors
D, I = index.search(xq, k) # search
print("search done, time cost:{}".format(time.time() - a ))
print(I[:5]) # neighbors of the 5 first queries
print(D[:5]) # distance of the return results
结果
generate data ..
generate data done, time cost:1.2140648365020752
start to build index..
build index done, time cost:1.621246337890625e-05
is trained: True
start to add query..
add data done, time cost:0.19340801239013672
1000000
start to search..
search done, time cost:8.8518226146698
[[ 774 175 917 413]
[ 98 66 357 596]
[ 504 263 740 738]
[ 756 1102 454 353]
[ 352 381 545 264]]
[[6.8088226 6.8297997 6.8576965 6.8842087]
[7.1985703 7.3808327 7.4176826 7.491226 ]
[7.2444 7.460396 7.4658546 7.6690903]
[6.293293 6.414982 6.531719 6.7257767]
[7.6404266 7.77285 7.8547783 8.039719 ]]
IndexIVFFlat
先聚类再索引。索引初始化后先在数据集上执行train方法（本质是在对数据集做kmeans聚类），再add数据。查询时先查询与各簇中心距离，找到最近的nprobe个簇后在相应的簇中遍历向量，找到最终结果。
from test_dataset import *
import faiss
"""
先聚类再搜索，可以加快检索速度
先将xb中的数据进行聚类（聚类的数目是超参），nlist: 聚类的数目
nprobe: 在多少个聚类中进行搜索，默认为1, nprobe越大，结果越精确，但是速度越慢
"""
nlist = 100 #聚类的数目
k = 4
a = time.time()
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist)
print("build index done, time cost:{}".format(time.time() - a ))
a = time.time()
index.train(xb) # IndexIVFFlat是需要训练的，这边是在聚类
print("train index done, time cost:{}".format(time.time() - a ))
print("is trained: {}".format(index.is_trained))
a = time.time()
index.add(xb) # add may be a bit slower as well
print("add data done, time cost:{}".format(time.time() - a ))
a = time.time()
D, I = index.search(xq, k) # actual search
print("nprobe=1 search done, time cost:{}".format(time.time() - a ))
print(I[-5:]) # neighbors of the 5 last queries
index.nprobe = 10 # default nprobe is 1, try a few more
a = time.time()
D, I = index.search(xq, k)
print("nprobe=10 search done, time cost:{}".format(time.time() - a ))
print(I[-5:]) # neighbors of the 5 last queries
结果
generate data ..
generate data done, time cost:1.2480273246765137
build index done, time cost:5.1021575927734375e-05
train index done, time cost:0.19619107246398926
is trained: True
add data done, time cost:0.30707430839538574
nprobe=1 search done, time cost:0.18898820877075195
[[10179 10714 9864 8813]
[ 9762 10013 9720 10405]
[ 9229 9510 10039 10602]
[10584 9725 9787 9816]
[ 9464 9635 10199 9180]]
nprobe=10 search done, time cost:1.5632450580596924
[[10179 10714 9864 8813]
[ 9762 10013 9720 10405]
[ 9229 9510 10039 10602]
[10584 9725 9787 9816]
[ 9464 9635 10199 9180]]
可以看出建立索引后搜索速度显著提升，从1个类中搜索仅耗时不到0.19s，从10个类中搜索耗时1.56s，远小于暴力搜索。
IndexIVFPQ
基于Product Quantizer的压缩算法将向量压缩并存储。此方法可以在存储时压缩向量的大小，故搜索的时候也属于近似搜索。优点是可以减少存储空间。
from test_dataset import *
import faiss
"""
基于乘积量化（product quantizers）对存储向量进行压缩，节省存储空间
m：乘积量化中，将原来的向量维度平均分成多少份，d必须为m的整数倍
bits: 每个子向量用多少个bits表示
和之前输出的最优结果略有不同，牺牲精度，但是节约空间
"""
nlist = 100
m = 8 # number of subquantizers
bits = 8
k = 4
a = time.time()
quantizer = faiss.IndexFlatL2(d) # this remains the same
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, bits) # 8 specifies that each sub-vector is encoded as 8 bits
print("build index done, time cost:{}".format(time.time() - a ))
a = time.time()
index.train(xb)
print("train index done, time cost:{}".format(time.time() - a ))
print("is trained: {}".format(index.is_trained)) # IndexFlatL2不需要训练过程
a = time.time()
index.add(xb)
print("add data done, time cost:{}".format(time.time() - a ))
index.nprobe = 10 # make comparable with experiment above
a = time.time()
D, I = index.search(xq, k) # search
print("nprobe=10 search done, time cost:{}".format(time.time() - a ))
print(I[:5])
结果
generate data ..
generate data done, time cost:1.3283205032348633
build index done, time cost:0.005830287933349609
train index done, time cost:3.3282480239868164
is trained: True
add data done, time cost:0.8176021575927734
nprobe=10 search done, time cost:0.596820592880249
[[ 175 365 774 66]
[ 56 375 3 286]
[ 740 738 394 519]
[ 353 454 1350 233]
[ 87 262 47 642]]
索引选择指南
Faiss提供了多种索引，涵盖了各种侧重点，在准确度、内存容量、数据集大小等方面可以综合考虑并选择适合自己业务的index。官方对此提供了较为详细的教程，参见Guidelines-to-choose-an-index
分布式检索
Faiss本身集成了一个rpc库（实际为创建socket连接进行通信，实现参见代码），可以支持分布式检索，代码参加官方demo。其原理为创建多个线程（线程数与远程服务器数相同）分别创建连接请求远程计算服务，并收集计算结果。
由于我手头只有两台机器，所以参考官方demo进行了简单修改，分别将计算分布在两台机器上和一台机器上进行对比测试
# 分布式版本
dis_machine_ports = [
('172.16.160.21', 12010),
('localhost', 12011),
]
# 单机版本
single_machine_ports = [
('localhost', 12014),
]
两台机器配置均为32核64G，测试结果如下
running:distribution
index size: 1000000
recall@1: 0.990
time cost: 3.421029567718506
running:single
index size: 1000000
recall@1: 0.989
time cost: 5.748946666717529
可以看到分布式明显快于单机版，且召回率略有提升。
运行过程中发现分布式版本的两台机器CPU利用率不足1000%，而单机版的CPU使用超过2000%，推测在单机版CPU已成为瓶颈。
在同一台机器上起了两个端口进行对比测试：
# 单机多实例版本
dis_machine_ports = [
('localhost', 12010),
('localhost', 12011),
]
# 单机版本
single_machine_ports = [
('localhost', 12014),
]
测试结果如下
running:distribution
index size: 1000000
recall@1: 0.992
time cost: 5.340034246444702
running:single
index size: 1000000
recall@1: 0.989
time cost: 4.62516450881958
单机版快于分布式版，证明确实存在CPU瓶颈问题。
总结
faiss是一款非常优秀的向量检索库，具有出色的性能和丰富的向量索引、量化、检索方法可供选择，可以高效地解决向量的knn（ann）问题。但由于其缺少完整的整套解决方案（如横向扩容、高可用、监控等），还远不能称为系统，更不能直接用于大规模线上服务。如需线上使用，需在方案上整体设计，将faiss作为计算节点，兼具高可用性、可运维性、易扩展性、可观测性，打造可靠系统。
参考文档
Faiss的github
Getting-started
常见问题总结
【关于 Faiss 】 那些你不知道的事
相似向量检索库-Faiss-简介及原理
Faiss-facebook开源框架入门教程