在推荐系统的召回阶段,如Youtube DNN和DSSM双塔模型,向量的最邻近检索是必不可少的一步。
一般的做法不会让模型在线预测召回,而是先离线将向量存储,然后在线上进行向量的最邻近检索,作为模型的召回。
例如:离线训练模型后,将item向量存储至某种数据库,然后线上推理时,模型实时计算输出user向量,然后通过Annoy或Faiss进行内积的最邻近检索。
这篇文章将介绍两个常用的向量最邻近检索工具:Annoy和Faiss。
Annoy
安装
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的
支持的距离度量和索引
支持的距离度量:
下面详细的讲下几种常用的索引类型:
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机器发送查询请求,然后合并查询结果。
-
大家想更自己实验性能和准确率的话,可以使用公开数据集ANN_SIFT