python怎么控制速度_如何加速faiss python search速度？

Faiss是Facebook提出用来处理大规模向量搜索召回问题的一种工具，在embedding大行其道的今天，有关向量之间计算相似度的场景很多，faiss是该类场景的首选工具。

线上serving时，使用faiss进行向量搜索召回，有一个大的向量集合，向量的量级一般是几十w，然后通过faiss构建索引；在服务进行请求时，一般是一个或者一组去请求faiss，寻找最邻近的向量，由于faiss底层是c++实现，并且每次服务请求的查询量少，所以serving的时候一般耗时不会有太大问题。

但是有一些离线大批量查询的场景，比如通过bert embedding好两部分商品的title，然后用一类商品标题的向量去另一类商品中查找最相似的商品。这类场景下，待索引的商品集合A一般很大，像我遇到的有60w、构建索引的商品集合B有20w，向量都是稠密向量，维度d=768维，这样一个faiss search场景耗时巨大。怎么优化？低阶索引+单条查询

首先，为了数据保密性，构造一组数据，用来模拟当时的数据；d=768 表示向量的维度；

nb = 200000 表示构建索引的向量size

nq = 5000 表示search向量size

import numpy

d = 768

nb = 200000

nq = 5000

np.random.seed(520)

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.

这样，数据就构造好了，其中xb是用来build faiss的，xq是用来search faiss的。

首先我们用低阶索引+单条查询，单条查询是每次search只查一条的意思，这是耗时最长的一种方案。

import faiss

index = faiss.IndexFlatL2(d)

index.add(xb)

import time

t1 = time.time()

for i in range(5000):

D, I = index.search(xb[i*1:(i+1)*1], 1)

t2 = time.time()

print '总计耗时:', t2-t1

print '单条耗时:', (t2-t1)/5000.0总计耗时: 293.038972139s

单条耗时: 0.0586077944279s低阶索引+batch查询

上面可以看到一条一条去查效率是很慢的，faiss是支持batch查询的，所以我们分批用batch去查

import time

t1 = time.time()

for i in range(500):

D, I = index.search(xb[i*10:(i+1)*10], 1)

t2 = time.time()

print '总计耗时:', t2-t1

print '单条耗时:', (t2-t1)/5000.0总计耗时: 48.9216761589

单条耗时: 0.00978433523178

由上面可以看到，把one search 改成 batch search即可加快search的速度，并且加速可观。那是不是batch越大越好？batch极限是多少？

可以看到当batch=10000及以上，单条耗时已经不再降低，说明已经到达上限了，这里是和机器性能有关；one search时cpu 使用率是100%，batch search时cpu使用率是100%~200%，不会都跑满。

初级加速faiss python search速度： one search 改成 batch search高阶索引+batch查询

这里只使用IndexIVFFlat高阶索引，不做其他高阶索引的对比。使用高阶索引是一种“trade precision for speed”，以精度换速度的方式。所以在使用高阶索引时，要做到精度和速度的平衡，这里是有参数可调的。

IndexIVFFlat是将向量分配到不同的单元格内，每个单元格有自己的中心点，原理上课k-means类似，在search可以控制每次search的单元格个数，这种索引是需要训练的。有两个关键的参数nlist和nprobe，nlist代表划分成多少个单元格，nprobe代表每次search时在几个单元格search。当nprobe=nlist时，就等效于低阶索引。

nlist = 2500

quantizer = faiss.IndexFlatL2(d)

index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)

index.train(xb)

index.add(xb)

高阶索引具体耗时对比如下：

可以看到，使用高阶索引后，one search 时，比低阶索引快75倍；相同bach search时，高阶索引比低阶索引快很多。

但是要指出的时，在使用高阶索引时，有一个index.nprobe = 10控制每次搜索的网格单元数，nprobe越小，速度越快，精度越差；nprobe越大，速度越慢，精度越好；我在使用时，nprobe设定为10，nlist=2500，精度几乎和暴力搜索一致，没什么精度损失，这两个参数需要自己根据各自的数据调节。

在调参nlist 和 nprobe时，可以构建一份低阶索引和一份高阶索引，使用同一份向量去search，默认低阶索引暴力搜索是正确的作为对照，diff高阶索引召回的向量，就可以计算出精度损失，然后继续调参即可。

进阶加速faiss python search速度：高阶索引+batch search