大模型应用--向量数据库FAISS

向量数据库Faiss是Facebook AI研究院开发的一种高效的相似性搜索和聚类的库。它能够快速处理大规模数据,并且支持在高维空间中进行相似性搜索。Faiss的工作,就是把我们自己的候选向量集封装成一个index数据库,它可以加速我们检索相似向量Top K的过程,一些最有用的算法是在 GPU 上实现的。

安装

pip install faiss-cpu
pip install faiss-gpu

通过上面的命令能够将faiss安装到环境中。

Faiss的基本使用

Faiss检索相似向量TopK的工程基本都能分为三步:

  1. 得到向量库;
  2. 用faiss 构建index,并将向量添加到index中;
  3. 用faiss index 检索。
import numpy as np
import faiss
d = 128                            # dimension
nb = 10000                         # database size
np.random.seed(1234)             # make reproducible
xb = np.random.random((nb, d)).astype('float32')
index = faiss.IndexFlatL2(d)   # build the index
print(index.is_trained)
index.add(xb)                  # add vectors to the index
print(index.ntotal)

nq = 5                          # number of query vectors
k = 4                           # we want 4 similar vectors
Xq = np.random.random((nq, d)).astype('float32')
D, I = index.search(Xq, k)     # sanity check

I是一个数组,它包含了每个查询向量的最近的4个向量的索引。D是一个数组,它包含了这些向量的距离。Faiss的强大之处在于它可以处理任何可以表示为向量的数据,包括图片和文件。

Faiss常用index算法

Faiss的索引是需要被训练的,下面介绍官方实现的几种基础索引类型。

聚类

ncentroids = 1024 
niter = 20 
verbose = True 
d = x.shape [1] 
kmeans = faiss.Kmeans (d, ncentroids, niter=niter, verbose=verbose) 
kmeans.train (x)

结果中心点存储在.kmeans.centroids 中。

PCA

# random training data 
mt = np.random.rand(1000, 40).astype('float32')
mat = faiss.PCAMatrix (40, 10)
mat.train(mt)
assert mat.is_trained
tr = mat.apply(mt)
# print this to show that the magnitude of tr's columns is decreasing
print (tr ** 2).sum(0)

量化器

支持的量化器类型有:

  • ScalarQuantizer:分别在线性范围内量化每个向量分量。
  • ProductQuantizer:对子向量执行向量量化
  • AdditiveQuantizer:将向量编码为码书条目的总和,详细信息请参见 Addtive Quantizers。 可以以多种方式训练加法量化器,因此有子类ResidualQuantizer ,LocalSearchQuantizer, ProductAdditiveQuantizer.

具体索引类型

暴力检索

  • 优点:该方法是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中添加向量

IVFx Flat :倒排暴力检索

  • 优点:IVF主要利用倒排的思想,在文档检索场景下的倒排技术是指,一个key-word后面挂上很多个包含该词的doc,由于key-word数量远远小于doc,因此会大大减少了检索的时间。在向量中如何使用倒排呢?可以拿出每个聚类中心下的向量ID,每个中心ID后面挂上一堆非中心向量,每次查询向量的时候找到最近的几个中心ID,分别搜索这几个中心下的非中心向量。通过减小搜索范围,提升搜索效率。
  • 缺点:速度也还不是很快。
  • 适用情况:相比Flat会大大增加检索的速度,建议百万级别向量可以使用。通常可以采用 4096, 或者 16384 个聚类中心的 IVF_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)

PQx :乘积量化

  • 优点:利用乘积量化的方法,改进了普通检索,将一个向量的维度切成x段,每段分别进行检索,每段向量的检索结果取交集后得出最后的Top-K。因此速度很快,而且占用内存较小,召回率也相对较高。
  • 缺点:召回率相较于暴力检索,下降较多。
  • 适用情况:内存及其稀缺,并且需要较快的检索速度,不那么在意召回率。
  • 参数:PQx中的x为将向量切分的段数,因此,x需要能被向量维度整除,且x越大,切分越细致,时间复杂度越高。
dim, measure = 64, faiss.METRIC_L2 
param =  'PQ16' 
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          
index.train(xb)                                 
index.add(xb)

IVFxPQy 倒排乘积量化

  • 优点:工业界大量使用此方法,各项指标都均可以接受,利用乘积量化的方法,改进了IVF的k-means,将一个向量的维度切成x段,每段分别进行k-means再检索。
  • 缺点:集百家之长,自然也集百家之短
  • 适用情况:一般来说,各方面没啥特殊的极端要求的话,最推荐使用该方法!
    参数:IVFx,PQy,其中的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)                    

LSH 局部敏感哈希

  • 原理:哈希对大家再熟悉不过,向量也可以采用哈希来加速查找,我们这里说的哈希指的是局部敏感哈希(Locality Sensitive Hashing,LSH),不同于传统哈希尽量不产生碰撞,局部敏感哈希依赖碰撞来查找近邻。高维空间的两点若距离很近,那么设计一种哈希函数对这两点进行哈希计算后分桶,使得他们哈希分桶值有很大的概率是一样的,若两点之间的距离较远,则他们哈希分桶值相同的概率会很小。
  • 优点:训练非常快,支持分批导入,index占内存很小,检索也比较快。
  • 缺点:召回率非常拉垮。
  • 使用情况:候选向量库非常大,离线检索,内存资源比较稀缺的情况。
dim, measure = 64, faiss.METRIC_L2  
param =  'LSH'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为True
index.add(xb)

HNSWx、IndexHNSWFlat

  • 这是一种基于图检索的改进方法,HNSWx中的x为构建图时每个点最多连接多少个节点,x越大,构图越复杂,查询越精确,当然构建index时间也就越慢,x取4~64中的任何一个整数。
  • 优点:不需要训练,基于图检索的改进方法,检索速度极快,10亿级别秒出检索结果,而且召回率几乎可以媲美Flat,能达到惊人的97%。检索的时间复杂度为O(loglogn),几乎可以无视候选向量的量级了。并且支持分批导入,极其适合线上任务,毫秒级别实时反应时间。
  • 缺点:构建索引极慢,占用内存极大(是Faiss中最大的,大于原向量占用的内存大小);添加数据不支持指定数据ID,不支持从索引中删除数据。
  • 适用情况:不在乎内存,并且有充裕的时间来构建index。
dim, measure = 64, faiss.METRIC_L2   
param =  'HNSW64' 
index = faiss.index_factory(dim, param, measure)  
print(index.is_trained)                          # 此时输出为True 
index.add(xb)
Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×