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

一，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

公众号算法美食屋后台回复关键词：源码，获取本文notebook源代码。