相似度检索Faiss模型
Posted xiaopihaierletian
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了相似度检索Faiss模型相关的知识,希望对你有一定的参考价值。
1. faiss作用
相似度检索TopK的问题一般的解决方案是暴力检索,循环遍历所有向量计算相似度然后得出TopK,但是当向量数量巨大时,这种方法及其耗时,Faiss的出现就很好地解决了这个问题。
2. faiss介绍
Faiss的全称是Facebook AI Similarity Search是FaceBook的AI团队针对大规模相似度检索问题开发的一个工具,使用C++编写,有python接口,对10亿量级的索引可以做到毫秒级检索的性能。Faiss的工作,就是把我们自己的候选向量集封装成一个index数据库,它可以加速我们检索相似向量TopK的过程,其中有些索引还支持GPU构建,可谓是强上加强。
3. QUICK START
FAISS的使用一般包括如下三格步骤:
step1: 构建向量库,一般可以直接通过词向量的平均或者通过BERT等预训练模型直接拿到一个句子的向量,来构建向量库
import numpy as np
d = 64 # 向量维度
nb = 100000 # index向量库的数据量
nq = 10000 # 待检索query的数目
np.random.seed(1234)
xb = np.random.random((nb, d)).astype('float32') # index向量库的向量
xq = np.random.random((nq, d)).astype('float32') # 待检索的query向量
step2: 构建index,并将向量添加到index中。这里我们选用暴力检索的方法FlatL2,L2代表构建的index采用的相似度度量方法为L2范数,即欧氏距离
import faiss
index = faiss.IndexFlatL2(d) # 创建索引时必须指定向量的维度d
print(index.is_trained) # 输出为True,代表该类index不需要训练,只需要add向量进去即可
index.add(xb) # 将向量库中的向量加入到index中
print(index.ntotal) # 输出index中包含的向量总数,为100000
step3: 检索TopK相似query
k = 4 # topK的K值
D, I = index.search(xq, k)# xq为待检索向量矩阵,返回的I为每个待检索query最相似TopK的索引list,D为其对应的距离
print(I[:5])
print(D[:5])
得到如下输出,第一个矩阵是索引,第一列0,1,2,3,4表示待检索矩阵自己的索引,从第二距离矩阵也可以看出,自己和自己的距离为0
[[ 0 393 363 78]
[ 1 555 277 364]
[ 2 304 101 13]
[ 3 173 18 182]
[ 4 288 370 531]]
[[ 0. 7.17517328 7.2076292 7.25116253]
[ 0. 6.32356453 6.6845808 6.79994535]
[ 0. 5.79640865 6.39173603 7.28151226]
[ 0. 7.27790546 7.52798653 7.66284657]
[ 0. 6.76380348 7.29512024 7.36881447]]
tips 1
在实际使用时,构建索引一般都使用faiss.index_factory 方法,基本所有的index都支持这种构建索引方法,上面步骤2中构建索引可以变成如下方式:
dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)
dim: 指定向量的维度
param: 是传入index的参数,代表需要构建什么类型的索引
measure:度量方法,目前支持两种,欧氏距离和inner product,即内积。因此,要计算余弦相似度,只需要将向量归一化后,使用内积度量即可。参数为faiss.METRIC_INNER_PRODUCT
tips 2
一些索引可以保存整型的ID,每个向量可以指定一个ID,当查询相似向量时,会返回相似向量的ID及相似度(或距离)。如果不指定,将按照添加的顺序从0开始累加。其中IndexFlatL2不支持指定ID。IndexFlatL2不支持指定id,但是可以通过IDMAP的方式实现,如下代码
ids = [2,10, 100,...]
ids = np.array(ids)
index = faiss.index_factory(768, "IDMap, Flat")
index.add_with_ids(save_embedding, ids) # 指定id,save_embedding为向量库
下面方法与上面类似
index = faiss.IndexFlatL2(d)
ids = np.arange(100000, 200000)
index2 = faiss.IndexIDMap(index)
index2.add_with_ids(xb, ids)
4. Faiss常用index优缺点及使用场景
4.1 Flat :暴力检索
优点:该方法是Faiss所有index中最准确的,召回率最高的方法,没有之一;
缺点:速度慢,占内存大
使用情况:向量候选集很少,在50万以内,并且内存不紧张
构建方法:
dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)
index.is_trained # 输出为True
index.add(xb) # 向index中添加向量
4.2 IVFx Flat :倒排暴力检索
优点:该方法是利用倒排技术对暴力检索加速,速度快于Flat
缺点:还不是很快,并且检索召回率下降
使用情况:同Flat,
参数:IVFx中的x是k-means聚类中心的个数
构建方法:
dim, measure = 64, faiss.METRIC_L2
param = 'IVF100, Flat' # 代表k-means聚类中心为100,
index = faiss.index_factory(dim, param, measure)
print(index.is_trained) # 此时输出为False,因为倒排索引需要训练k-means,
index.train(xb) # 因此需要先训练index,再add向量
index.add(xb)
4.3 HNSWx
优点:该方法为基于图检索的改进方法,检索速度极快,10亿级别秒出检索结果,而且召回率几乎可以媲美Flat,能达到惊人的97%。检索的时间复杂度为loglogn,几乎可以无视候选向量的量级了。并且支持分批导入,极其适合线上任务,毫秒级别体验。
缺点:构建索引极慢,占用内存极大(是Faiss中最大的,大于原向量占用的内存大小)
参数:HNSWx中的x为构建图时每个点最多连接多少个节点,x越大,构图越复杂,查询越精确,当然构建index时间也就越慢,x取4~64中的任何一个整数。
使用情况:不在乎内存,并且有充裕的时间来构建index
构建方法:
dim, measure = 64, faiss.METRIC_L2
param = 'HNSW64'
index = faiss.index_factory(dim, param, measure)
print(index.is_trained) # 此时输出为True
index.add(xb)
4.4 PQx :乘积量化
优点:利用乘积量化的方法,改进了普通k-means,将一个向量的维度切成x段,每段分别进行k-means。因此速度很快,而且占用内存较小,召回率也相对较高
缺点:召回率相较于暴力检索,下降较多。
使用情况:内存及其稀缺,并且需要较快的检索速度,不那么在意召回率
参数:PQx中的x为将向量切分的段数,因此,x需要能被向量维度整除,且x越大,切分越细致,时间复杂度越高
构建方法:
dim, measure = 64, faiss.METRIC_L2
param = 'PQ16'
index = faiss.index_factory(dim, param, measure)
print(index.is_trained) # 此时输出为False,因为倒排索引需要训练k-means,
index.train(xb) # 因此需要先训练index,再add向量
index.add(xb)
4.5 IVFxPQy 倒排乘积量化
优点:工业界大量使用此方法,各项指标都均可以接受。
缺点:集百家之长,自然也集百家之短
使用情况:同PQx
参数:IVFxPQy,其中的x和y同上
构建方法:
dim, measure = 64, faiss.METRIC_L2
param = 'IVF100, PQ16'
index = faiss.index_factory(dim, param, measure)
print(index.is_trained) # 此时输出为False,因为倒排索引需要训练k-means,
index.train(xb) # 因此需要先训练index,再add向量 index.add(xb)
更多参考文献[3]
5. 使用GPU
参考Running on GPUs
5.1 使用单个gpu
res = faiss.StandardGpuResources() # 声明gpu资源
# 构建一个 flat (CPU) 索引
index_flat = faiss.IndexFlatL2(d)
# 将cpu索引加入到gpu中
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
# 接下来的操作与一般情况类似
gpu_index_flat.add(xb) # add vectors to the index
print(gpu_index_flat.ntotal)
k = 4 # we want to see 4 nearest neighbors
D, I = gpu_index_flat.search(xq, k) # actual search
print(I[:5]) # neighbors of the 5 first queries
print(I[-5:]) # neighbors of the 5 last queries
5.2 使用多个GPU
ngpus = faiss.get_num_gpus()
print("number of GPUs:", ngpus)
cpu_index = faiss.IndexFlatL2(d)
gpu_index = faiss.index_cpu_to_all_gpus(cpu_index) # build the index
gpu_index.add(xb) # add vectors to the index
print(gpu_index.ntotal)
k = 4 # we want to see 4 nearest neighbors
D, I = gpu_index.search(xq, k) # actual search
print(I[:5]) # neighbors of the 5 first queries
print(I[-5:]) # neighbors of the 5 last queries
参考
faiss简介及示例
Faiss流程与原理分析
Faiss入门及应用经验记录
以上是关于相似度检索Faiss模型的主要内容,如果未能解决你的问题,请参考以下文章