10分钟入门faiss相似向量检索
共 4672字，需浏览 10分钟
·
2023-08-18 21:21
一，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)，适合千万上亿规模以及更大规模的图索引，缺点是构建索引过程较慢，占用很大的存储。
公众号算法美食屋后台回复关键词：源码，获取本文notebook源代码。
二，Faiss安装
pip install faiss-cpu # faiss-gpu (一般来说cpu足够快了)
三，代码范例
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)
输出如下：
============================== 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