Faiss入门及应用经验记录

转载自：https://zhuanlan.zhihu.com/p/357414033

 

导语

说起相似度检索TopK的问题，相信很多算法er在实际工程中会经常遇到，对此我们一般的解决方案是暴力检索，循环遍历所有向量计算相似度然后得出TopK。但当向量数量级达到百万千万甚至上亿级别，这时候你再用暴力检索就会显得很呆 ... ...

Faiss的出现就很好地解决了这个问题，笔者总结了在工程中使用Faiss的一些经验，记录下给需要的童鞋（语言为Python，因为本菜鸡不会C++）。动动小手给点个赞呗。

---

1. 什么是Faiss？

Faiss的全称是Facebook AI Similarity Search，是FaceBook的AI团队针对大规模相似度检索问题开发的一个工具，使用C++编写，有python接口，对10亿量级的索引可以做到毫秒级检索的性能。

简单来说，Faiss的工作，就是把我们自己的候选向量集封装成一个index数据库，它可以加速我们检索相似向量TopK的过程，其中有些索引还支持GPU构建，可谓是强上加强。

---

2. Faiss简单上手

首先，Faiss检索相似向量TopK的工程基本都能分为三步：

1. 得到向量库；
2. 用faiss 构建index，并将向量添加到index中；
3. 用faiss index 检索。

好吧......这貌似和废话没啥区别，参考把大象装冰箱需要几个步骤。本段代码摘自Faiss官方文档，很清晰，基本所有的index构建流程都遵循这个步骤。

• 第一步，得到向量：

import numpy as np
d = 64                                           # 向量维度
nb = 100000                                      # index向量库的数据量
nq = 10000                                       # 待检索query的数目
np.random.seed(1234)             
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.                # index向量库的向量
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.                # 待检索的query向量

• 第二步，构建索引，这里我们选用暴力检索的方法FlatL2，L2代表构建的index采用的相似度度量方法为L2范数，即欧氏距离：

import faiss          
index = faiss.IndexFlatL2(d)             
print(index.is_trained)         # 输出为True，代表该类index不需要训练，只需要add向量进去即可
index.add(xb)                   # 将向量库中的向量加入到index中
print(index.ntotal)             # 输出index中包含的向量总数，为100000 

• 第三步，检索TopK相似query：

k = 4                     # topK的K值
D, I = index.search(xq, k)# xq为待检索向量，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离
print(I[:5])
print(D[-5:])

• 打印输出为：

>>> 
[[  0 393 363  78] 
 [  1 555 277 364] 
 [  2 304 101  13] 
 [  3 173  18 182] 
 [  4 288 370 531]]  
[[ 0.          7.17517328  7.2076292   7.25116253]  
 [ 0.          6.32356453  6.6845808   6.79994535]  
 [ 0.          5.79640865  6.39173603  7.28151226]  
 [ 0.          7.27790546  7.52798653  7.66284657]  
 [ 0.          6.76380348  7.29512024  7.36881447]]

3. Faiss常用index优缺点及使用场景

Faiss之所以能加速，是因为它用的检索方式并非精确检索，而是模糊检索。既然是模糊检索，那么必定有所损失，我们用召回率来表示模糊检索相对于精确检索的损失。

在我们实际的工程中，候选向量的数量级、index所占内存的大小、检索所需时间（是离线检索还是在线检索）、index构建时间、检索的召回率等都是我们选择index时常常需要考虑的地方。

首先，我建议关于Faiss的所有索引的构建，都统一使用faiss.index_factory，基本所有的index都支持这种构建索引方法。

以第二章的代码为例，构建index方法和传参方法建议修改为：

dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)

• 三个参数中，dim为向量维数；
• 最重要的是param参数，它是传入index的参数，代表需要构建什么类型的索引；
• measure为度量方法，目前支持两种，欧氏距离和inner product，即内积。因此，要计算余弦相似度，只需要将vecs归一化后，使用内积度量即可。

3.1 Flat ：暴力检索

• 优点：该方法是Faiss所有index中最准确的，召回率最高的方法，没有之一；
• 缺点：速度慢，占内存大。
• 使用情况：向量候选集很少，在50万以内，并且内存不紧张。
• 构建方法：

dim, measure = 64, faiss.METRIC_L2
param =  'Flat'
index = faiss.index_factory(dim, param, measure)
index.is_trained                                   # 输出为True
index.add(xb)                                      # 向index中添加向量

3.2 IVFx Flat ：倒排暴力检索

• 优点：该方法是利用倒排技术对暴力检索加速，速度快于Flat
• 缺点：还不是很快，并且检索召回率下降
• 使用情况：同Flat，
• 参数：IVFx中的x是k-means聚类中心的个数
• 构建方法：

dim, measure = 64, faiss.METRIC_L2 
param =  'IVF100, Flat'                          # 代表k-means聚类中心为100,   
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)                                     

3.3 PQx ：乘积量化

• 优点：利用乘积量化的方法，改进了普通k-means，将一个向量的维度切成x段，每段分别进行k-means。因此速度很快，而且占用内存较小，召回率也相对较高
• 缺点：召回率相较于暴力检索，下降较多。
• 使用情况：内存及其稀缺，并且需要较快的检索速度，不那么在意召回率
• 参数：PQx中的x为将向量切分的段数，因此，x需要能被向量维度整除，且x越大，切分越细致，时间复杂度越高
• 构建方法：

dim, measure = 64, faiss.METRIC_L2 
param =  'PQ16' 
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)          

3.4 IVFxPQy 倒排乘积量化

• 优点：工业界大量使用此方法，各项指标都均可以接受。
• 缺点：集百家之长，自然也集百家之短
• 使用情况：同PQx
• 参数：IVFxPQy，其中的x和y同上
• 构建方法：

dim, measure = 64, faiss.METRIC_L2  
param =  'IVF100, PQ16'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means， 
index.train(xb)                                  # 因此需要先训练index，再add向量 index.add(xb)       

3.5 LSH 局部敏感哈希

• 优点：训练非常快，支持分批导入，index占内存很小，检索也比较快
• 缺点：召回率非常拉垮。在候选语料比较多的时候（百万级别），检索也不是特别快，大概是秒级别的，不太适用于线上部署 ... ...
• 使用情况：候选向量库非常大，离线检索，内存资源比较稀缺的情况
• 构建方法：

dim, measure = 64, faiss.METRIC_L2  
param =  'LSH'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为True
index.add(xb)       

3.6 HNSWx （最重要的放在最后说）

• 优点：该方法为基于图检索的改进方法，检索速度极快，10亿级别秒出检索结果，而且召回率几乎可以媲美Flat，能达到惊人的97%。检索的时间复杂度为loglogn，几乎可以无视候选向量的量级了。并且支持分批导入，极其适合线上任务，毫秒级别体验。
• 缺点：构建索引极慢，占用内存极大（是Faiss中最大的，大于原向量占用的内存大小）
• 参数：HNSWx中的x为构建图时每个点最多连接多少个节点，x越大，构图越复杂，查询越精确，当然构建index时间也就越慢，x取4~64中的任何一个整数。
• 使用情况：不在乎内存，并且有充裕的时间来构建index
• 构建方法：

dim, measure = 64, faiss.METRIC_L2   
param =  'HNSW64' 
index = faiss.index_factory(dim, param, measure)  
print(index.is_trained)                          # 此时输出为True 
index.add(xb)

4. 使用Faiss时踩过的坑和一些经验

4.1 以下是我的同事对不同index做过的一些对比实验结果：

召回率、搜索时间对比

召回率、检索时间、内存消耗、构建时间对比

4.2 Faiss可以组合传参

Faiss内部支持先将向量PCA降维后再构建index，param设置如下：

param = 'PCA32, IVF100, PQ16'

代表将向量先降维成32维，再用IVF100 PQ16的方法构建索引。

同理可以使用：

param = 'PCA32, HNSW32'

可以用来处理HNSW内存占用过大的问题。

4.3 Faiss在构建索引时，有时生成的vecs会很大，向index中添加的向量很有可能无法一次性放入内存中，怎么办呢？

• 这时候，索引的可分批导入index的性质就起了大作用了；
• 如何来知道一种index是否可以分批add呢？一般来说在未对index进行train操作前，如果一个index.is_trained = True，那么它就是可以分批add的；
• 如果是index.is_trained = False，就不能分批导入，当然，其实强行对index.is_trained = False的索引分批add向量是不会报错的，只不过内部构建索引的逻辑和模型都在第一个batch数据的基础上构建的，比如PCA降维，其实是拿第一个batch训练了一个PCA模型，后续add进入的新向量都用这个模型降维，这样会导致后续batch的向量失真，影响精度；
• 由上可得，只要一个index的param中包含PCA，那他就不能分批add；
• 可以分批导入的index为：HNSW、Flat、LSH。

4.4 如果我们的需求，既想PCA降维减小index占用内存，还想分批add向量，该怎么办？

可以使用sklean中的增量pca方法，先把数据降维成低维的，再将降维后的向量分批add进索引中，增量pca使用方法和pca一致：

from sklearn.decomposition import IncrementalPCA

4.5 关于HNSW

• HNSW虽好，可不要贪杯哦，这里坑还是蛮多的
• HNSW的构建过程有时很短，有时却又很长，500万量级的向量快时可以15分钟构建完毕，慢则有可能花费两小时，这和HNSW构建多层图结构时的生成函数有关。
• HNSW在使用以下这种构建索引方式时，有一个bug，这个bug会导致如果measure选用内积的方法度量，但最后构建的索引的度量方式仍然是欧氏距离：

index = faiss.index_factory(dim, param, measure)  

如果想构建以内积（余弦相似度）为基准的HNSW索引，可以这样构建：

index = faiss.IndexHNSWFlat(dim, x,measure)  # measure 选为内积，x为4~64之间的整数

• HNSW占用内存真的很大，500万条768维的向量构建的索引占用内存17G，而同样数据下LSH仅占用500M，emmm所以自行体会吧。

4.6 Faiss所有的index仅支持浮点数为float32格式

Faiss仅支持浮点数为np.float32格式，其余一律不支持，所以用Faiss前需要将向量数据转化为float32，否则会报错！这也告诉大家，想用降低精度来实现降低index内存占用是不可能的！

5.Faiss构建索引时，如何权衡数据量级和内存占用来选择index呢？

这里有一份Faiss的官方文档，里面写的很清楚：

Faiss官网，如何选择index类型​github.com

关于Faiss中index的更多详细信息，请参考官网：

Faiss index factory介绍​github.com