faiss的python接口使用

1. 简介

faiss是一种ann（Approximate Nearest Neighbor）库，可以用于特征的入库，检索。

不仅可以在cpu上使用，还可以利用GPU进行检索，提高检索的速度。

具体可以参考：https://github.com/facebookresearch/faiss

2. 安装

cpu版本，适用于各个系统

pip install faiss-cpu

cpu + gpu版本，目前不适用于windows系统

pip install faiss-gpu

3. 示例

1. 新建索引

   import faiss
   
   # 传入特征维度
   dim = 2048
   
   # IndexFlatIP表示利用内积来比较特征的相似度
   # 这里一般会让提取的特征进行L2归一化，那么内积就等于余弦相似度
   index_ip = faiss.IndexFlatIP(dim)
   
   # IndexFlatL2表示利用L2距离来比较特征的相似度
   index_l2 = faiss.IndexFlatL2(dim)
   

2. 添加特征入库

   import numpy as np
   
   # 新建一个特征，维度为2048，shape为(1, 2048)
   feature = np.random.random((1, 2048)).astype('float32')
   index_ip.add(feature)
   
   # 当然，也可以一次性添加多个特征
   features = np.random.random((10, 2048)).astype('float32')
   index_ip.add(features)
   
   # 打印index_ip包含的特征数量
   print(index_ip.ntotal) 
   

3. 自己指定每个特征的id

   在第2步中，添加特征的id是根据特征入库的顺序对应的，如果想自己指定id，可以用IndexIDMap包装一层，代码如下所示：

   index_ids = faiss.IndexFlatIP(2048)
   index_ids = faiss.IndexIDMap(index_ids)
   
   # 添加特征，并指定id，注意添加的id类型为int64
   ids = 20
   feature_ids = np.random.random((1, 2048)).astype('float32')
   index_ids.add_with_ids(feature_ids, np.array((ids,)).astype('int64'))
   

   这里要注意，包装的索引必须是空的，即一开始新建索引之后，就进行包装，不能入库过一些特征后中途再包装。

4. 检索

   feature_search = np.random.random((1, 2048)).astype('float32')
   
   # 检索最相似的topK个特征
   topK = 5
   D, I = index_ip.search(feature_search, topK)
   # 返回的D表示相似度（或者距离）, I表示检索的topK个特征id（索引）
   

5. 保存/读取索引文件

   faiss的另一个优点是，可以将保存着特征的索引持久化，保存为文件，类似数据库，这样就不用每次都提取特征了。

   # 保存索引
   faiss.write_index(index_ip, 'my.index')
   
   # 读取索引
   index = faiss.read_index('my.index')
   

6. GPU的使用

   # 利用单个gpu
   res = faiss.StandardGpuResources()
   
   gpu_index = faiss.index_cpu_to_gpu(res, 0, index_ip)
   

其他操作可以参考faiss在github上的地址。

结束。