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

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

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

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

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

Annoy

Annoy Github

安装

pip install annoy

支持的距离度量

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

• 欧式距离euclidean

• 内积dot

  

• 汉明距离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维向量下：

具体的性能和召回准确测试见官方的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_Lp
• METRIC_Canberra
• METRIC_BrayCurtis
• METRIC_JensenShannon

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

• IndexFlatL2：欧式距离的精确搜索，即要求100%召回；
• IndexFlatIP：同上，针对内积
• IndexIVFFlat：将所有向量通过聚类的方法划分至nlist个空间，然后在搜索的时候，比较目标向量与所有空间的中心距离，选出nprobe个最近的空间，然后比较这些空间里面的向量。此方法适用于保证较高的召回率下实现高效的查询性能，同时适用于欧式距离和内积
• IndexScalarQuantizer：向量压缩，如将向量的FLOAT(4个字节)压缩至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维向量下：

单位: s	IndexFlatL2	IndexIVFFlat	IndexFlatL2-GPU	IndexIVFFlat-GPU
batch_size=1	0.12	0.004	0.013	<0.001
batch_size=10	0.77	0.004	0.06	~=0.001
batch_size=100	1.13	0.024	0.05	0.002
batch_size=1000	6.3	0.17	0.21	0.009

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

总结

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

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

3. 如果想要分布式部署faiss或annoy的话，一般的思路是：

   在不同的salve机器上，分别存储等量的数据，并创建索引；

   然后，在额外一台master机器，分别想salve机器发送查询请求，然后合并查询结果。

   （Distributed index，demo_client_server）

4. 大家想更自己实验性能和准确率的话，可以使用公开数据集ANN_SIFT