10分钟入门faiss相似向量检索

2023-09-05 13:41:37 浏览数 (1)

一,faiss简介

faiss全称 Facebook AI Similarity Search,是FaceBook的AI团队针对大规模向量 进行 TopK 相似向量 检索 的一个工具,使用C 编写,有python接口,对10亿量级的索引可以做到毫秒级检索的性能

使用faiss 搭配合适的model和embedding函数,可以帮助我们 构建 人脸识别,相似图片检索,LLM知识库问答,推荐系统召回模块 等应用。

faiss的主要原理是构建base vectors向量数据的index索引,然后利用索引对search vectors 实现 TopK 相似向量检索。

faiss支持许多不同的构建索引的方式,以下是一些较推荐使用的类型。

1,Flat:暴力精确检索,全局最优,适合数十万级。

2,IVF100,Flat:倒排暴力检索(100聚类后暴力检索),非全局最优但召回高,适合数百万级。

3, HNSW64: 图网络检索,Hierarchical NSW(Navigable Small World),每个节点64条边,检索复杂度log(logn),适合千万上亿规模以及更大规模的图索引,缺点是构建索引过程较慢,占用很大的存储。

二,Faiss安装

代码语言:javascript复制
pip install faiss-cpu # faiss-gpu (一般来说cpu足够快了)

三,代码范例

代码语言:javascript复制
import faiss
import numpy as np 

# 〇,基本参数设置
d = 64                                           # 向量维度
nb = 100000                                      # index向量库的数据量
nq = 1000                                        # 待检索query的数目
index_type = 'Flat'                              # index 类型
metric_type = faiss.METRIC_INNER_PRODUCT         # 度量(相似度/距离)类型


# 一,准备向量库向量
print('============================== 1,base vector ==============================')
np.random.seed(1234)             
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0]  = np.arange(nb) / 1000.                # index向量库的向量
faiss.normalize_L2(xb)
print('xb.shape = ',xb.shape,'n')


# 二,准备查询向量
print('============================== 2,query vector ==============================')
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0]  = np.arange(nq) / 1000.                # 待检索的query向量
faiss.normalize_L2(xq)
print('xq.shape = ',xq.shape,'n')

# 三,构建向量库索引
print('============================== 3,create&train ==============================')
index = faiss.index_factory(d,index_type,metric_type)    #等价于 faiss.IndexFlatIP(d)     
print('index.is_trained=',index.is_trained)    # 输出为True,代表该类index不需要训练,只需要add向量进去即可
index.train(xb)
index.add(xb)                                      # 将向量库中的向量加入到index中
print('index.ntotal=',index.ntotal,'n')           # 输出index中包含的向量总数,为100000 


# 四,相似向量查询
print('============================== 4, search ==============================')
k = 4                       # topK的K值
D, I = index.search(xq, k)  # xq为待检索向量,返回的I为每个待检索query最相似TopK的索引list,D为其对应的距离

print('nearest vector ids:n',I[:5],'n')
print('metric(distances/scores) to query:n',D[-5:],'n')

# 五,增删索引向量
print('============================== 5, add&remove ==============================')
xa = np.random.random((10000, d)).astype('float32')
xa[:, 0]  = np.arange(len(xa)) / 1000.                
faiss.normalize_L2(xa)
index.add(xa)
print('after add, index.ntotal=',index.ntotal) 
index.remove_ids(np.arange(1000,1111))
print('after remove, index.ntotal=',index.ntotal,'n') 

# 六,保存加载索引
print('============================== 6, write&read ==============================')
faiss.write_index(index, "large.index")
index_loaded = faiss.read_index('large.index')
print('index_loaded.ntotal=', index_loaded.ntotal)

输出如下:

代码语言:javascript复制
============================== 1,base vector ==============================
xb.shape =  (100000, 64) 

============================== 2,query vector ==============================
xq.shape =  (1000, 64) 

============================== 3,create&train ==============================
index.is_trained= True
index.ntotal= 100000 

============================== 4, search ==============================
nearest vector ids:
 [[ 207  381 1394 1019]
 [ 300  911  142  526]
 [ 838 1541  527  148]
 [ 196  359  184  466]
 [ 526  120  917  765]] 

metric(distances/scores) to query:
 [[0.87687665 0.86128205 0.85667723 0.85451   ]
 [0.8702938  0.86668813 0.85934925 0.8523142 ]
 [0.862915   0.85807455 0.85384977 0.8499449 ]
 [0.8692     0.86600477 0.8647547  0.8634621 ]
 [0.8539625  0.84914947 0.84744585 0.8432568 ]] 

============================== 5, add&remove ==============================
after add, index.ntotal= 110000
after remove, index.ntotal= 109889 

============================== 6, write&read ==============================
index_loaded.ntotal= 109889

参考文章: 1,《Faiss入门以及应用经验记录》 https://zhuanlan.zhihu.com/p/357414033

2,《ANN召回算法之HNSW》 https://zhuanlan.zhihu.com/p/379372268

0 人点赞