Faiss(Facebook开源的高效相似搜索库)学习小记

1. 写在前面

faiss是在设计推荐系统入门竞赛之新闻推荐中学习到的一个非常好用的工具包，这个是Facebook AI团队开源的针对聚类和相似性搜索库，为稠密向量提供高效相似度搜索和聚类，支持十亿级别向量的搜索，是目前最为成熟的近似近邻搜索库。它包含多种搜索任意大小向量集（备注：向量集大小由RAM内存决定）的算法，以及用于算法评估和参数调整的支持代码。Faiss用C++编写，并提供与Numpy完美衔接的Python接口。除此以外，对一些核心算法提供了GPU实现。

当时的应用场景就是面对20万用户点击过的所有历史文章，大概是3万多篇历史文章，我们针对这每一篇历史文章，都得基于embedding相似度，从36万多篇新闻中召回100篇与其相似的文章，这个当时还不知道faiss的时候，这个运算算了3天也没有跑出结果来。这种暴力搜索的方式无疑是非常花费时间的，也就是在这个场合下，学习到了Faiss这个包，大大提高了搜索效率，把3天还没有搞定的问题几个小时就搞定了。下面记录一下这个包的基本使用过程，参考了官方给出的文档https://github.com/facebookresearch/faiss/wiki/Getting-started

faiss工具包一般使用在推荐系统中的向量召回部分。在做向量召回的时候要么是u2u,u2i或者i2i，这里的u和i指的是user和item.我们知道在实际的场景中user和item的数量都是海量的，我们最容易想到的基于向量相似度的召回就是使用两层循环遍历user列表或者item列表计算两个向量的相似度，但是这样做在面对海量数据是不切实际的，faiss就是用来加速计算某个查询向量最相似的topk个索引向量。

faiss查询的原理：

faiss使用了PCA和PQ(Product quantization乘积量化)两种技术进行向量压缩和编码，当然还使用了其他的技术进行优化，但是PCA和PQ是其中最核心部分。

• 主成分分析（PCA）原理总结
• 实例理解product quantization算法

2. 安装

faiss的安装网上有一大堆教程，我这里只记录一下我的安装方式，我这里使用了最方便快速的方式，即基于Anaconda进行安装的，facebook推出了faiss的conda安装包，在conda安装时会自行安装所需的libgcc, mkl, numpy模块。

有两个版本：faiss的cpu版本目前仅支持Linux和MacOS操作系统，gpu版本提供可在Linux操作系统下用CUDA8.0/CUDA9.0/CUDA9.1编译的版本。
注意，上面语句中的cuda90并不会执行安装CUDA的操作，需要提前自行安装。

我安装的CPU版本， 代码如下：

#安装cpu版本
#更新conda
conda update conda
#先安装mkl
conda install mkl
#安装faiss-cpu
conda install faiss-cpu -c pytorch
#测试安装是否成功
python -c "import faiss"

安装GPU版本， 这个我没试过，因为anaconda环境有点问题，怕麻烦就没有尝试。应该也是一样的：

#安装gpu版本
#确保已经安装了CUDA，否则会自动安装cup版本。
conda install faiss-gpu -c pytorch # 默认 For CUDA8.0
conda install faiss-gpu cuda90 -c pytorch # For CUDA9.0
conda install faiss-gpu cuda91 -c pytorch # For CUDA9.1

当然不嫌麻烦，好像也可以编译安装， 但时间宝贵，不要浪费在这种环境上哈哈。 我没试，具体可以上网查查。

基本使用

Faiss处理大规模d维向量近邻检索的问题，Faiss中所有向量以行矩阵的形式储存和使用，下面我们通过一个实例看看基本的使用。

实例中我们用xb表示所有待索引的向量集合 (也就是上面我说的36万篇文章的embedding)，xq表示查询向量集合(上面的用户看过的3万篇文章的embedding)，nb和nq分别表示xb、xq集合中向量数量。 这个例子里面的待索引向量个数是100000， 查询向量个数是10000

1. 构建待检索向量和查询向量

   import numpy as np
   # 随机种子确定
   np.random.seed(1234)
   
   # 向量维度
   d = 64
   
   # 待索引向量size
   nb = 100000
   # 查询向量size
   nq = 10000
   
   
   #为了使随机产生的向量有较大区别进行人工调整向量  产生两类向量
   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.
   

2. 建立索引并添加向量
   Faiss通过Index对象进行向量的封装与预处理，Faiss提供了很多种索引类型，具体如下：
   
   

   我们首先test暴力搜索精准L2距离搜索，对应的索引对象为IndexFlatL2。所有向量在建立前需要明确向量的维度d，大多数的索引还需要训练阶段来分析向量的分布。但是，对于L2暴力搜索来说没有训练的必要。
   

   import faiss
   
   # 建立索引
   index = faiss.IndexFlatL2(d)
   
   print(index.is_trained)    # True
   

   Index对象训练好之后，对于Index有两个操作供调用，分别为add和search。

   • add方法用于向Index中添加xb向量
   • search方法用于在add向量后的索引中检索xq的若干近邻。

   
   Index还有两个状态变量is_trainedbool类型，用于指示index是否已被训练)和ntotal(指示索引的数量)。此外，index还有IDs添加的方法。

   # 索引中添加向量
   index.add(xb)
   print(index.ntotal)     # 100000
   

3. 近邻搜索
   通过Index检索xq中的数据，faiss支持批量数据检索，通过search方法返回的检索结果包括两个矩阵，分别为近邻向量的索引序号和xq中元素与近邻的距离大小。

   # 返回每个查询向量的近邻个数
   k = 4
   
   # 检索check
   D, I = index.search(xb[:5], k)
   print(I)
   print(D)
   
   #xq检索结果
   D, I = index.search(xq, k)
   # 前五个检索结果展示
   print(I[:5])
   # 最后五个检索结果展示
   print(I[-5:])
   

   结果如下：
   
   上面所示，IndexFlatL2是暴力检索的索引，为了加速检索speed，可以使用faiss的IndexIVFFlat索引对象。IndexIVFFlat的使用需要进行训练阶段，并需要指定其他索引作为量化器，与检索相关的参数为nlist和nprobe。
   IndexIVFFlat索引先利用粗量化器将检索向量划分到Voronoi单元中并建立倒排索引，检索阶段将根据输入向量和probe参数定位到对应的Voronoi cell中进行近邻搜素。具体使用方法如下：

4. IndexIVFlat检索

   nlist = 100
   k = 4
   
   # 量化器索引
   quantizer = faiss.IndexFlatL2(d)
   
   # 指定用L2距离进行搜索，若不指定默认为內积
   index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
   
   # 索引训练
   assert not index.is_trained
   index.train(xb)   
   
   assert index.is_trained
   # 向量添加
   index.add(xb)
   
   # 检索
   D, I = index.search(xq, k)
   
   # 最后五个检索结果
   print(I[-5:])
   
   # 多探针检索
   index.nprobe = 10
   #最后五个检索结果
   D, I = index.search(xq, k)
   print(I[-5:])
   

目前用到的就这些知识，详细的可以去查阅官方文档了，如果再遇到新的，再进行更新。

这里补充在基本项目里面的使用， 一般faiss用在向量召回阶段， 给定一个原item向量， 期待从海量item中得到与原item相似的TopK， 这里记录一个真实使用的代码。 后面遇到faiss，基本上就是这个处理逻辑拿过来就可以:

def embedding_recall(df, item_emb_df, department, subject, topK=50):

    # 筛选query_df
    query_df = df.loc[(df['department']==department) & (df['subject']==subject)].reset_index(drop=True)
    item_emb_df = item_emb_df.loc[(item_emb_df['department']==department) & (item_emb_df['subject']==subject)].reset_index(drop=True)

    # 题目索引与题目id的字典映射
    item_idx_2_rawid_dict = dict(zip(item_emb_df.index, item_emb_df['item_id']))

    item_emb_np = np.array(list(map(eval, item_emb_df['content_embedding'].values))).astype('float32')   # 这里的类型必须是float32

    item_index = faiss.IndexFlatIP(item_emb_np.shape[1])
    item_index.add(item_emb_np)

    # query_df 建立索引对应关系
    query_idx_2_rawid_dict = dict(zip(query_df.index, query_df['item_id']))
    query_emb_np = np.array(list(map(eval, query_df['content_embedding'].values))).astype('float32')

    # 相似度查询，给每个索引位置上的向量返回topk个item以及相似度
    sim, idx = item_index.search(query_emb_np, topK+1)

    # 将向量检索的结果保存成原始id的对应关系
    item_sim_dict = collections.defaultdict(dict)

    for target_idx, sim_value_list, rele_idx_list in zip(range(query_emb_np.shape[0]), sim, idx):
        target_raw_id = query_idx_2_rawid_dict[target_idx]

        # 从1开始是为了去掉题目本身
        for rele_idx, sim_value in zip(rele_idx_list[1:], sim_value_list[1:]):
            rele_raw_id = item_idx_2_rawid_dict[rele_idx]
            item_sim_dict[target_raw_id][rele_raw_id] = item_sim_dict.get(target_raw_id, {}).get(rele_raw_id, 0) + sim_value
    
    return item_sim_dict

这里的item_emb_df，就是我们的海量item以及向量， 传入的格式是DataFrame， 两列， 第一列是item_id, 第二列是content_embedding。

当然， 我这里是根据实际情况， 先根据一些条件进行一波筛选， 也就是海量item， 先经过条件筛选一波，就比如有各种类的文章， 我们召回相似的时候， 先限制同类里面去找是一个道理。

筛选环节可以在传入之前， 也可以在传入之后的第一步， 这个看情况，但下面的就是正常使用faiss的逻辑。

1. 建立索引->item映射， 这个是为了，当faiss返回结果的时候，也是返回的索引->embedding， 所以我们得通过这个映射得到item->embedding
2. 把item_emb_np和query_emb_np存储成二维数组的形式， 且格式必须是float32， 不能是float64.
3. 建立faiss_index， 把item_emb_np通过add方法加入
4. 通过search方法搜索topK
5. 根据映射拿到最终结果返回

其中用到faiss的核心代码就3句:

item_index = faiss.IndexFlatIP(item_emb_np.shape[1])
item_index.add(item_emb_np)

# 相似度查询，给每个索引位置上的向量返回topk个item以及相似度
sim, idx = item_index.search(query_emb_np, topK+1)