原文链接:https://blog.csdn.net/kanbuqinghuanyizhang/java/article/details/80774609

简介

faiss是为稠密向量提供高效相似度搜索和聚类的框架。由Facebook AI Research研发。 具有以下特性。

1、提供多种检索方法
2、速度快
3、可存在内存和磁盘中
4、C++实现,提供Python封装调用。
5、大部分算法支持GPU实现
下面给出一些快速链接方便查找更多内容。

github
官方文档
Troubleshooting
官方安装文档

安装

仅允许使用conda安装,不可以用pip,亲测pip会报错

  1. # 更新conda
  2. conda update conda
  3. # 先安装mkl
  4. conda install mkl
  5. # faiss提供gpu和cpu版,根据服务选择
  6. # cpu版本
  7. conda install faiss-cpu -c pytorch
  8. # gpu版本 -- 记得根据自己安装的cuda版本安装对应的faiss版本,不然会出异常。使用命令:nvcc -V 查看
  9. conda install faiss-gpu cudatoolkit=8.0 -c pytorch # For CUDA8
  10. conda install faiss-gpu cudatoolkit=9.0 -c pytorch # For CUDA9
  11. conda install faiss-gpu cudatoolkit=10.0 -c pytorch # For CUDA10
  12. # 校验是否安装成功
  13. python -c "import faiss"
  14. # 可能会有mkl的报错,尤其是执行index.search,导入以下环境即可
  15. export LD_PRELOAD=/<anacondas3安装路径>/anaconda3/lib/libmkl_rt.so

Quick Start

这里先给出官方提供的demo来感受一下faiss的使用。
首先构建训练数据和测试数据

  1. import numpy as np
  2. d = 64 # dimension
  3. nb = 100000 # database size
  4. nq = 10000 # nb of queries
  5. np.random.seed(1234) # make reproducible
  6. xb = np.random.random((nb, d)).astype('float32')
  7. xb[:, 0] += np.arange(nb) / 1000.
  8. xq = np.random.random((nq, d)).astype('float32')
  9. xq[:, 0] += np.arange(nq) / 1000.

上面我们构建了shape为[100000,64]的训练数据xb,和shape为[10000,64]的查询数据xq。
然后创建索引(Index)。faiss创建索引对向量预处理,提高查询效率。
faiss提供多种索引方法,这里选择最简单的暴力检索L2距离的索引:IndexFlatL2。
创建索引时必须指定向量的维度d。大部分索引需要训练的步骤。IndexFlatL2跳过这一步。
当索引创建好并训练(如果需要)之后,我们就可以执行add和search方法了。add方法一般添加训练时的样本,search就是寻找相似相似向量了。
一些索引可以保存整型的ID,每个向量可以指定一个ID,当查询相似向量时,会返回相似向量的ID及相似度(或距离)。如果不指定,将按照添加的顺序从0开始累加。其中IndexFlatL2不支持指定ID。

  1. import faiss # make faiss available
  2. index = faiss.IndexFlatL2(d) # build the index
  3. print(index.is_trained)
  4. index.add(xb) # add vectors to the index
  5. print(index.ntotal)

我们有了包含向量的索引后,就可以传入搜索向量查找相似向量了。

  1. k = 4 # we want to see 4 nearest neighbors
  2. D, I = index.search(xq, k) # actual search
  3. print(I[:5]) # neighbors of the 5 first queries
  4. print(D[-5:]) # neighbors of the 5 last queries

上面代码中,我们定义返回每个需要查询向量的最近4个向量。查询返回两个numpy array对象D和I。D表示与相似向量的距离(distance),维度,I表示相似用户的ID。

我们可以得到类似于下面的结果

  1. [[ 0 393 363 78]
  2. [ 1 555 277 364]
  3. [ 2 304 101 13]
  4. [ 3 173 18 182]
  5. [ 4 288 370 531]]
  6. [[ 0. 7.17517328 7.2076292 7.25116253]
  7. [ 0. 6.32356453 6.6845808 6.79994535]
  8. [ 0. 5.79640865 6.39173603 7.28151226]
  9. [ 0. 7.27790546 7.52798653 7.66284657]
  10. [ 0. 6.76380348 7.29512024 7.36881447]]

加速搜索

如果需要存储的向量太多,通过暴力搜索索引IndexFlatL2速度很慢,这里介绍一种加速搜索的方法的索引IndexIVFFlat。翻译过来叫倒排文件,其实是使用K-means建立聚类中心,然后通过查询最近的聚类中心,然后比较聚类中的所有向量得到相似的向量。
创建IndexIVFFlat时需要指定一个其他的索引作为量化器(quantizer)来计算距离或相似度。
这里同使用IndexFlatL2对比,在add方法之前需要先训练。

下面简述示例中的几个参数。
faiss.METRIC_L2: faiss定义了两种衡量相似度的方法(metrics),分别为faiss.METRIC_L2、faiss.METRIC_INNER_PRODUCT。一个是欧式距离,一个是向量内积。
nlist:聚类中心的个数
k:查找最相似的k个向量
index.nprobe:查找聚类中心的个数,默认为1个。

代码示例如下

  1. nlist = 100 #聚类中心的个数
  2. k = 4
  3. quantizer = faiss.IndexFlatL2(d) # the other index
  4. index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
  5. # here we specify METRIC_L2, by default it performs inner-product search
  6. assert not index.is_trained
  7. index.train(xb)
  8. assert index.is_trained
  9. index.add(xb) # add may be a bit slower as well
  10. D, I = index.search(xq, k) # actual search
  11. print(I[-5:]) # neighbors of the 5 last queries
  12. index.nprobe = 10 # default nprobe is 1, try a few more
  13. D, I = index.search(xq, k)
  14. print(I[-5:]) # neighbors of the 5 last queries

减少内存

2018-02-22之后版本添加了磁盘存储inverted indexes的方式,使用可参考demo.

上面我们看到的索引IndexFlatL2和IndexIVFFlat都会全量存储所有的向量在内存中,为满足大的数据量的需求,faiss提供一种基于Product Quantizer(乘积量化)的压缩算法编码向量大小到指定的字节数。此时,存储的向量时压缩过的,查询的距离也是近似的。关于乘积量化的算法可自行搜索。

下面给出demo。类似IndexIVFFlat,这里使用的是IndexIVFPQ

  1. nlist = 100
  2. m = 8 # number of bytes per vector
  3. k = 4
  4. quantizer = faiss.IndexFlatL2(d) # this remains the same
  5. index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)
  6. # 8 specifies that each sub-vector is encoded as 8 bits
  7. index.train(xb)
  8. index.add(xb)
  9. D, I = index.search(xb[:5], k) # sanity check
  10. print(I)
  11. print(D)
  12. index.nprobe = 10 # make comparable with experiment above
  13. D, I = index.search(xq, k) # search
  14. print(I[-5:])

之前我们定义的维度为d = 64,向量的数据类型为float32。这里压缩成了8个字节。所以压缩比率为 (64*32/8) / 8 = 32
返回的结果如下,第一个向量同自己的距离为1.40704751,不是0。因为如上所述返回的是近似距离,但是整体上返回的最相似的top k的向量ID没有变化。

  1. [[ 0 608 220 228]
  2. [ 1 1063 277 617]
  3. [ 2 46 114 304]
  4. [ 3 791 527 316]
  5. [ 4 159 288 393]]
  6. [[ 1.40704751 6.19361687 6.34912491 6.35771513]
  7. [ 1.49901485 5.66632462 5.94188499 6.29570007]
  8. [ 1.63260388 6.04126883 6.18447495 6.26815748]
  9. [ 1.5356375 6.33165455 6.64519501 6.86594009]
  10. [ 1.46203303 6.5022912 6.62621975 6.63154221]]

简化索引的表达

通过上面IndexIVFFlat和IndexIVFPQ我们可以看到,他们的构造需要先提供另外一个index。类似的,faiss还提供pca、lsh等方法,有时候他们会组合使用。这样组合的对构造索引会比较麻烦,faiss提供了通过字符串表达的方式构造索引。
如,下面表达式就能表示上面的创建IndexIVFPQ的实例。

  1. index = faiss.index_factory(d, "IVF100,PQ8")

这里有一点文档中没有提到的,通过查看c++代码,index_factory方法还有第三个参数,就是上面说的metric。可传入的就上面两种。

  1. Index *index_factory (int d, const char *description_in, MetricType metric)

更多的组合实例可以看demo
每类索引的简写可查询Basic indexes

GPU使用

注意有些索引不支持GPU,哪些支持哪些不支持可查询Basic indexes
可通过faiss.get_num_gpus()查询有多少个gpu

  1. ngpus = faiss.get_num_gpus()
  2. print("number of GPUs:", ngpus)

使用gpu的完整示例。

1、使用一块gpu

  1. # build a flat (CPU) index
  2. index_flat = faiss.IndexFlatL2(d)
  3. # make it into a gpu index
  4. gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)

2、使用全部gpu

  1. cpu_index = faiss.IndexFlatL2(d)
  2. gpu_index = faiss.index_cpu_to_all_gpus(cpu_index) # build the index
  3. gpu_index.add(xb) # add vectors to the index
  4. print(gpu_index.ntotal)
  5. k = 4 # we want to see 4 nearest neighbors
  6. D, I = gpu_index.search(xq, k) # actual search
  7. print(I[:5]) # neighbors of the 5 first queries
  8. print(I[-5:]) # neighbors of the 5 last queries