推荐系统的向量检索工具: 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_LinfMETRIC_LpMETRIC_CanberraMETRIC_BrayCurtisMETRIC_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
总结
- Annoy不支持批量插入和查询，仅支持一种索引类型，单步查询速度快；
- Faiss支持批量插入和查询，支持100%召回、量化压缩等多种索引类型，并且能够使用GPU加速，适用：
- 内存资源不足时，使用向量量化压缩的索引；
- 具备GPU资源；
- 需要batch查询（batch查询均摊下来比单个查询有很大的优势）；
-
如果想要分布式部署faiss或annoy的话，一般的思路是：
在不同的salve机器上，分别存储等量的数据，并创建索引；
然后，在额外一台master机器，分别想salve机器发送查询请求，然后合并查询结果。
（Distributed index，demo_client_server）
-
大家想更自己实验性能和准确率的话，可以使用公开数据集ANN_SIFT