基于 brpc + faiss 的矢量检索框架

基于 brpc + faiss 的矢量检索框架

基于 brpc + faiss 的矢量检索框架

概述

矢量检索，顾名思义，就是通过固定维度的浮点数矢量，在数据库里查找与之距离最近的 top-k 个矢量的过程，在深度学习领域是一个非常常见，而且非常实用的基础组件，在视频检索，文本召回，人脸识别验证，推荐系统等问题均有很多应用。

这里，我们提供一个阉割版的矢量检索服务框架，性能并不高，主要用于教学目的，不过可以基于此框架升级到比较实用的状态。

阉割版目前的特性包括：

支持并发的增删改查过程，查询 QPS 是上万级采用一种特殊的随机过程对分桶函数进行 online 训练，所以可以支持实时 + 并发的入库，以 online 方式构建索引，自动索引下，写入 QPS 约为 7000-8000支持简单的持久化方案提供一些简单的并发测试脚本单机上限差不多是亿级，不过对系统资源占用极大，阉割版没有提供量化，以及多索引方案，对系统资源占用也没有限制，容易出现问题，不建议加载这种量级数据。

性能评测 (测试机器为 32 核 cpu 机器，内存 60 G)：

200 维 x 10 M 数据，写入 QPS 约为 6000，检索 QPS 约为 2000-3000128 维 x 1M 数据，写入 QPS 约为 8000，检索 QPS 约为 8000

启动服务

确保系统安装有 docker，并保证在代码的根目录下运行如下命令：

构建镜像

TARGET=brpc_faiss_server:v1docker build -t ${TARGET} .

启动镜像，比如我们设置端口为 8330，需要将 src/config.h 里面的 IP_PORT 修改为 8330，数据缓存位置设置为 tmp，如有需要可以换到你想要的任意目录下

SERVER_TMP_DIR=/absolute/path/to/your/tmp/dirmkdir ${SERVER_TMP_DIR}docker run -itd --runtime nvidia --name brpc_faiss_server -v ${SERVER_TMP_DIR}:/data/saved_rocksdb_faiss -v `pwd`:/opt/brpc_server -p 8330:8330 ${TARGET} /bin/bash

启动服务

docker exec brpc_faiss_server /bin/sh -c /opt/brpc_server/start.sh

使用服务

获取 docker 的宿主机 ip 地址 xx.xxx.x.xxx，也就是服务地址内部提供 protobuf 的接口，不过为了方便展示，这里用 http 方式调度python 可以直接使用 request 模块进行 http 方式调度服务，也可以使用 curl 命令访问服务监控面板，可以在浏览器打开 http://xx.xxx.x.xxx:8330/status ，进行监控，这里的 xx.xxx.x.xxx 就是宿主机 ip

提供了 6 个服务，每个服务的 url 为：

创建一个检索数据库： http://xx.xxx.x.xxx:8330/brpc_faiss_server/create批量添加数据：http://xx.xxx.x.xxx:8330/brpc_faiss_server/batch_add删除数据库：http://xx.xxx.x.xxx:8330/brpc_faiss_server/delete查询数据库状态：http://xx.xxx.x.xxx:8330/brpc_faiss_server/status从指定位置，恢复数据库：http://xx.xxx.x.xxx:8330/brpc_faiss_server/restore进行矢量检索：http://xx.xxx.x.xxx:8330/brpc_faiss_server/search

增加一个库：

import requestsurl = " http://xx.xxx.x.xxx:8330/brpc_faiss_server/create"request_data = { "db_name": "word.emb", "feature_dim": 200, "feature_version": "v1", "search_type": 0, "similarity_type": "L2", "search_device": "cpu",}request_data = json.dumps(request_data)result = requests.post(url=url, data=request_data)

查询服务状态：

import requestsurl = " http://xx.xxx.x.xxx:8330/brpc_faiss_server/status"request_data = { "db_name": "word.emb",}request_data = json.dumps(request_data)result = requests.post(url=url, data=request_data)

删除某个库：

import requestsurl = " http://xx.xxx.x.xxx:8330/brpc_faiss_server/delete"request_data = { "db_name": "word.emb"}request_data = json.dumps(request_data)result = requests.post(url=url, data=request_data)

批量添加：

import requestsurl = " http://xx.xxx.x.xxx:8330/brpc_faiss_server/batch_add"data = []for cid, fid, feature in list(zip(cids, fids, features)): item = {"cid": cid, "feature_id": fid, "b64_feature": feature} data.append(item)request_data = { "db_name": "word.emb", "feature_version": "v1", "feature_dim": 200, "data": data,}request_data = json.dumps(request_data)result = requests.post(url=url, data=request_data)

矢量检索：

import requestsurl = " http://xx.xxx.x.xxx:8330/brpc_faiss_server/search"request_data = { "db_name": "word.emb", "b64_feature": "NxxWPhZuKb4Iym29TfaPPjro0j2Yh8w9A7BBPfGC6D0drE8+6kB2P...", "topk": 10,}request_data = json.dumps(request_data)result = requests.post(url=url, data=request_data)

并发测试

这里只是提供简单的 python 接口做并发测试，机器： 32 核 cpu 机器

首先安装的 python 库，其中 Locust 库是并发压测工具，可以依靠单机多核，轻松形成几万的 QPS：

pip install requests, Locust, numpy

- bigann 1M 数据集：ftp://ftp.irisa.fr/local/texmex/corpus/sift.tar.gz，并解压到用户指定目录里面

cd 到 python 的 test_concurrent 目录里面

创建一个库，并执行并发入库：

./run.sh xx.xxx.x.xxx add.py 3m path_to_bigann

检索：

./run.sh xx.xxx.x.xxx search.py 3m path_to_bigann

在压测过程，在浏览器打开 http://xx.xxx.x.xxx:8330/status ，就可以查看性能面板，非常方便：

动态索引方案

定义 $k$ 个粗中心点 $c_i$，以及分桶函数 $p_c(i|x)$： $$ p_c(i|x) = \frac{e^{-\eta |x-c_i|^2}}{\sum_j e^{-\eta |x-c_j|^2}} $$ 给样本划分桶的时候，采取 $\arg\max_i p_c(i|x)$ 方式选择桶，当 $\eta$ 无穷大的时候，分桶函数就是最近邻分桶。

对于 faiss 来讲，分桶实现方式是构建 voronoi cell，方法是 kmean，每个 cell 都是一个桶，cell 内的样本都是桶内样本，不同 cell 之间的样本是桶间样本。然而，voronoi cell 本身有一定缺点：

有一定可能出现大量样本落在两个 cell 之间的边界上，导致检索需要跨桶判断边界样本，增加时间开销某个 cell 里面包含太多样本，导致检索某个桶的负担增加

我们提出，对于倒排矢量检索来讲，一个好的分桶策略应当满足如下要求：

每个桶内的样本距离桶中心的偏差都不大，数学表述为 $ \sum_x |x-c_i|^2$ 要小桶间样本的距离要拉大

这样，检索过程只要在单个桶里检索便可以，

给定一个数据分布 $p(x)$，我们希望每个样本距离自己最近的中心点更近，距离剩余中心点更远： $$ l(x)=\sum_i |x-c_i|^2 p(i|x) + \sum_i |x-c_i|^2 (1-p(i|x)) $$

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。