推荐系统的向量检索工具: Annoy & Faiss

╰+攻爆jí腚メ 2022-09-07 14:37 446阅读 0赞

在推荐系统的召回阶段，如Youtube DNN和DSSM双塔模型，向量的最邻近检索是必不可少的一步。

一般的做法不会让模型在线预测召回，而是先离线将向量存储，然后在线上进行向量的最邻近检索，作为模型的召回。

**例如：离线训练模型后，将item向量存储至某种数据库，然后线上推理时，模型实时计算输出user向量，然后通过Annoy或Faiss进行内积的最邻近检索。**

这篇文章将介绍两个常用的向量最邻近检索工具：**Annoy**和**Faiss**。

# Annoy #

[Annoy Github][]

## 安装 ##

pip install annoy

## 支持的距离度量 ##

**Annoy仅支持树结构的索引类型。**

*  `欧式距离euclidean`

![请添加图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NneXVhbnNoaQ_size_16_color_FFFFFF_t_70]

*  `内积dot`
    
    ![请添加图片描述][5095a26846674f80aee3bdc079404784.png]
 *  `汉明距离hamming`
    
    两个二进制字符串的距离，即可以用来计算0-1向量的距离，实际对应位置不同的数量，如：
    
    11011001 ⊕ 10011101 = 2
 *  `angular、manhattan`

## python sdk ##

*  创建ann实例
    
        import numpy as np
        from annoy import AnnoyIndex
        
        ann = AnnoyIndex(dim, "euclidean")
 *  插入数据
    
        dim = 256
        data = np.random.random([1000000, dim]).astype(np.float32)
        # 插入ID和对应的向量
        # 不支持批量插入
        for i, arr in enumerate(data):
        		ann.add_item(i, arr)
 *  构建索引
    
        n_trees = 512
        ann.build(n_trees)
 *  查询
    
        # 根据ID检索
        qid = 10
        res = get_nns_by_item(query, 100)
        
        # 根据向量检索
        query = np.random.random([dim])
        res = get_nns_by_vector(query, 100)
 *  保存和加载
    
        # 保存至磁盘
        ann.save("your path")
        
        # 从磁盘中加载到内存
        ann = AnnoyIndex(dim, "euclidean")
        ann.load("your path", prefault=True)

## 性能 ##

在1核的cpu，100W的128维向量下：

![请添加图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NneXVhbnNoaQ_size_16_color_FFFFFF_t_70 1]

具体的性能和召回准确测试见[官方的Benchmark][Benchmark]

# Faiss #

## 安装 ##

如果已经正确安装CUDA：

# cpu version
    pip install faiss-cpu
    
    # gpu version
    pip install faiss-gpu

如果以上安装出现问题的话，可以使用官方推荐的方法：

# CPU-only version
    $ conda install -c pytorch faiss-cpu
    
    # GPU(+CPU) version
    $ conda install -c pytorch faiss-gpu
    
    # or for a specific CUDA version
    $ conda install -c pytorch faiss-gpu cudatoolkit=10.2 # for CUDA 10.2

或者通过源码编译：https://github.com/facebookresearch/faiss/blob/master/INSTALL.md

**注意：GPU版本的faiss是包含CPU的**

## 支持的距离度量和索引 ##

**支持的距离度量：**

*  欧式距离L2、内积、L1距离
 *  [`METRIC_Linf`][METRIC_Linf]
 *  [`METRIC_Lp`][METRIC_Lp]
 *  [`METRIC_Canberra`][METRIC_Canberra]
 *  [`METRIC_BrayCurtis`][METRIC_BrayCurtis]
 *  [`METRIC_JensenShannon`][METRIC_JensenShannon]

下面详细的讲下几种常用的索引类型：

*  `IndexFlatL2`：欧式距离的精确搜索，**即要求100%召回**；
 *  `IndexFlatIP`：同上，针对内积
 *  `IndexIVFFlat`：将所有向量通过聚类的方法划分至**nlist**个空间，然后在搜索的时候，比较目标向量与所有空间的中心距离，选出**nprobe**个最近的空间，然后比较这些空间里面的向量。**此方法适用于保证较高的召回率下实现高效的查询性能，同时适用于欧式距离和内积**
 *  `IndexScalarQuantizer`：向量压缩，如将向量的FLOAT(4个字节)压缩至1个字节。**适用于内存资源缺乏的场景。**
 *  [其他索引类型][Link 1]

## python sdk ##

精准召回的IndexFlatL2

import numpy as np
    import faiss
    import time
    
    # 构建数据集
    dims = 128
    np.random.seed(2020)
    data = np.random.random([1000000, dims]).astype('float32')
    
    # 构建index并插入数据
    index = faiss.IndexFlatL2(dims)  # build the index
    # 如果指定ID，可以使用：add_with_ids
    index.add(data)  # add vectors to the index
    print(index.ntotal)
    
    # 查询
    queries = np.random.random([1, dims]).astype('float32')
    D, I = index.search(queries, k)  # D和I分别对应距离和ID

IndexIVFFlat：需要先对数据train，进行聚类，然后再插入数据

nlist = 512
    
    quantizer = faiss.IndexFlatL2(dims)  # the other index
    index = faiss.IndexIVFFlat(quantizer, dims, nlist)
    # train聚类
    index.train(data)
    # 插入数据
    index.add(data) 
    
    # 查询
    queries = np.random.random([1, dims]).astype('float32')
    D, I = index.search(queries, k)  # D和I分别对应距离和ID

使用GPU加速

res = faiss.StandardGpuResources()
    # build a flat (CPU) index
    index_flat = faiss.IndexFlatL2(dims)
    # make it into a gpu index
    gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
    # add vectors to the index
    gpu_index_flat.add(data)        
    print(gpu_index_flat.ntotal)
    
    # 查询
    queries = np.random.random([1, dims]).astype('float32')
    D, I = gpu_index_flat.search(queries, k)  # D和I分别对应距离和ID

## 性能 ##

在1核的cpu&Tesla P4，100W的128维向量下：

<table> 
 <thead> 
  <tr> 
   <th>单位: s</th> 
   <th>IndexFlatL2</th> 
   <th>IndexIVFFlat</th> 
   <th>IndexFlatL2-GPU</th> 
   <th>IndexIVFFlat-GPU</th> 
  </tr> 
 </thead> 
 <tbody> 
  <tr> 
   <td>batch_size=1</td> 
   <td>0.12</td> 
   <td>0.004</td> 
   <td>0.013</td> 
   <td>&lt;0.001</td> 
  </tr> 
  <tr> 
   <td>batch_size=10</td> 
   <td>0.77</td> 
   <td>0.004</td> 
   <td>0.06</td> 
   <td>~=0.001</td> 
  </tr> 
  <tr> 
   <td>batch_size=100</td> 
   <td>1.13</td> 
   <td>0.024</td> 
   <td>0.05</td> 
   <td>0.002</td> 
  </tr> 
  <tr> 
   <td>batch_size=1000</td> 
   <td>6.3</td> 
   <td>0.17</td> 
   <td>0.21</td> 
   <td>0.009</td> 
  </tr> 
 </tbody> 
</table>

具体的性能和召回准确测试见[官方的Benchmark][Benchmark 1]

# 总结 #

1.  Annoy不支持批量插入和查询，仅支持一种索引类型，单步查询速度快；
2.  Faiss支持批量插入和查询，支持100%召回、量化压缩等多种索引类型，并且能够使用GPU加速，适用：

*  内存资源不足时，使用向量量化压缩的索引；
 *  具备GPU资源；
 *  需要batch查询（batch查询均摊下来比单个查询有很大的优势）；

1.  如果想要分布式部署faiss或annoy的话，一般的思路是：
    
    在不同的salve机器上，分别存储等量的数据，并创建索引；
    
    然后，在额外一台master机器，分别想salve机器发送查询请求，然后合并查询结果。
    
    （[Distributed index][]，[**demo\_client\_server**][demo_client_server]）
2.  大家想更自己实验性能和准确率的话，可以使用公开数据集[ANN\_SIFT][ANN_SIFT]

[Annoy Github]: https://github.com/spotify/annoy
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NneXVhbnNoaQ_size_16_color_FFFFFF_t_70]: /images/20220829/8f75401c805c4775b3f902bb1556024c.png
[5095a26846674f80aee3bdc079404784.png]: /images/20220829/e7111058213b44faa3f5f6602e9233a4.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NneXVhbnNoaQ_size_16_color_FFFFFF_t_70 1]: /images/20220829/c77de81d2c244fc2a5758287c64843bd.png
[Benchmark]: https://github.com/erikbern/ann-benchmarks
[METRIC_Linf]: https://en.wikipedia.org/wiki/Chebyshev_distance
[METRIC_Lp]: https://en.wikipedia.org/wiki/Lp_space
[METRIC_Canberra]: https://en.wikipedia.org/wiki/Canberra_distance
[METRIC_BrayCurtis]: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.braycurtis.html
[METRIC_JensenShannon]: https://en.wikipedia.org/wiki/Jensen%E2%80%93Shannon_divergence
[Link 1]: https://github.com/facebookresearch/faiss/wiki/Faiss-indexes
[Benchmark 1]: https://github.com/facebookresearch/faiss/blob/master/benchs/README.md
[Distributed index]: https://github.com/facebookresearch/faiss/wiki/Indexes-that-do-not-fit-in-RAM
[demo_client_server]: https://github.com/facebookresearch/faiss/blob/master/demos/demo_client_server_ivf.py
[ANN_SIFT]: http://corpus-texmex.irisa.fr/