faiss包装与性能对比

faiss包装与性能对比

在NLP的应用中,经常需要用到对向量的搜索,如果向量的数量级非常大,比如1千万,甚至上亿条,普通的方式就满足不了生产需要了,falcebook开源的faiss框架能够解决“海量向量搜索”的问题。

在网上也有很多faiss的相关文章,大多都是介绍安装以及对官方demo的运行测试。到底faiss的速度如何,使用内存情况如何,如何把faiss包装成服务与项目结合,带着这些问题,笔者做了一些尝试,把这些过程做一下分享。

项目说明

本项目的源码地址为:https://github.com/xmxoxo/faiss_test

项目是为了做一个对比实验,来看一下faiss的使用及速度。
我们随机生成10万个768维的向量,加载到内存中,
然后分别用普通的暴力搜索和faiss搜索两种方式去搜索,对比搜索的平时用时。

整个实验过程的开发日志见: develop log.txt

faiss介绍

Faiss是Facebook Ai Research开发的一款稠密向量检索工具。
简单来说就是向量搜索工具。
引用Faiss Wiki上面的一段简介说明下特点:
1.Faiss是针对稠密向量进行相似性搜索和聚类的一个高效类库。
2. 它包含可搜索任意大小的向量集的算法,这些向量集的大小甚至都不适合RAM。
3. 它还包含用于评估和参数调整的支持代码。
4. Faiss用C ++编写,并且有python2与python3的封装代码。
5. 一些最有用的算法在GPU上有实现。
6. Faiss是由Facebook AI Research开发的。

faiss安装

faiss需要在linux下安装,这里的测试环境是:
ubuntu 16.04
python 3.6.6
使用conda安装faiss是便捷的方式。

# 更新conda
conda update conda
# 先安装mkl
conda install mkl

# faiss提供gpu和cpu版,根据服务选择
# cpu版本
conda install faiss-cpu -c pytorch

# gpu版本 -- 记得根据自己安装的cuda版本安装对应的faiss版本,不然会出异常

# For CUDA8
conda install faiss-gpu cudatoolkit=8.0 -c pytorch 

# For CUDA9
conda install faiss-gpu cudatoolkit=9.0 -c pytorch 

# For CUDA10
conda install faiss-gpu cudatoolkit=10.0 -c pytorch 

注意,如果安装不成功时,可以尝试去掉 命令行最后的-c pytorch

# 校验是否安装成功
python -c "import faiss"

普通暴力搜索

为了对比搜索的速度,先用普通的暴力搜索,就是一个个计算余弦相似度,保留TOP N。
在测试工具里增加了交互,运行后可以自行测试。

运行命令:

python VecSearch_force.py

10万条数据的运行结果:

root@ubuntu:/mnt/sda1/transdat/VecSearch# python VecSearch_Lib.py
===========大批量向量余弦相似度计算-[暴力版]===========
随机生成100000个向量,维度:768
正在创建搜索器...
添加用时:0.027474秒
Used Memory: 672.36328125 MB
-----------------单条查询测试-----------------
搜索结果: [0.78086126 0.77952558 0.77949381 0.77775592 0.77547914] [ 541  443 1472  370  209]
显示查询结果,并验证余弦相似度...
索引号:  541, 距离:0.780861
索引号:  443, 距离:0.779526
索引号: 1472, 距离:0.779494
索引号:  370, 距离:0.777756
索引号:  209, 距离:0.775479
-----------------批量查询测试-----------------
批量测试次数:100 次,请稍候...
总用时:93 秒, 平均用时:932.549202 毫秒

faiss向量搜索

faiss向量搜索测试工具中,对faiss向量搜索进行了包装,
后续会发布了一个faiss向量搜索通用服务端。

在命令行中可以设置各类测试参数,支持GPU;

具体帮助信息可以使用以下命令查看:

python VecSearch_faiss.py -h

帮助信息如下:

VecSearch_faiss.py: error: argument -h/--help: ignored explicit argument 'elp'
(base) root@ubuntu:/mnt/sda1/transdat/VecSearch# python VecSearch_faiss.py -h
usage: VecSearch_faiss.py [-h] [--total TOTAL] [--dim DIM]
                          [--test_times TEST_TIMES] [--top_n TOP_N]
                          [--gpu GPU]

faiss速度测试工具

optional arguments:
  -h, --help            show this help message and exit
  --total TOTAL         总数据量
  --dim DIM             向量维度
  --test_times TEST_TIMES
                        测试次数
  --top_n TOP_N         每次返回条数
  --gpu GPU             使用GPU,-1=不使用,0=使用第1个,>0=使用全部

10万条768维向量搜索性能:

root@ubuntu:/mnt/sda1/transdat/VecSearch# python VecSearch_faiss.py
=========大批量向量余弦相似度计算-[faiss版]==========
随机生成100000个向量,维度:768
正在创建搜索器...
(0, 99999)
创建用时:2.116421秒
Used Memory: 1326.91015625 MB
-----------------单条查询测试-----------------
显示查询结果,并验证余弦相似度...
索引号: 1058, 距离:114.878273, 余弦相似度:0.771127
索引号:  541, 距离:115.051292, 余弦相似度:0.780861
索引号:  370, 距离:115.715881, 余弦相似度:0.777756
索引号:  209, 距离:115.731331, 余弦相似度:0.775479
索引号: 1472, 距离:115.832954, 余弦相似度:0.779494
总用时:9毫秒
-----------------批量查询测试-----------------
正在批量测试1000次,每次返回Top 5,请稍候...
总用时:1415毫秒, 平均用时:1.415620毫秒

运行截图

运行过程中的一些截图,可以看git源码里的images目录。
这里上传一个暴力搜索与faiss cpu搜索的对比图:

在这里插入图片描述

结束语

通过整个对比试验,可以得到以下的结果:

  • 对faiss向量搜索进行了类包装,可方便地做成一个独立的搜索服务端;
  • faiss的搜索速度非常快,基本在几个毫秒;
  • 对于更大数据量的向量搜索,难点在于向量数据文件的存储与加载

关于faiss向量搜索的后续研究还会陆续推出,个别课题已经基本完成:

  • faiss向量搜索通用服务端;
  • 使用faiss向量搜索加载腾讯880万词向量
  • simBERT+faiss 实现语义搜索;

欢迎加入群讨论:
QQ群:BERT模型爱好者 891629665
扫描下面的二维码加入。
在这里插入图片描述

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值