相似度检索Faiss模型

🤗专注公众号 funNLPer白嫖有用的知识🤗

1. faiss作用

相似度检索TopK的问题一般的解决方案是暴力检索，循环遍历所有向量计算相似度然后得出TopK，但是当向量数量巨大时，这种方法及其耗时，Faiss的出现就很好地解决了这个问题。

2. faiss介绍

Faiss的全称是Facebook AI Similarity Search是FaceBook的AI团队针对大规模相似度检索问题开发的一个工具，使用C++编写，有python接口，对10亿量级的索引可以做到毫秒级检索的性能。Faiss的工作，就是把我们自己的候选向量集封装成一个index数据库，它可以加速我们检索相似向量TopK的过程，其中有些索引还支持GPU构建，可谓是强上加强。

3. QUICK START

FAISS的使用一般包括如下三格步骤：

• step1: 构建向量库，一般可以直接通过词向量的平均或者通过BERT等预训练模型直接拿到一个句子的向量，来构建向量库

import numpy as np
d = 64                                           # 向量维度
nb = 100000                                      # index向量库的数据量
nq = 10000                                       # 待检索query的数目
np.random.seed(1234)             
xb = np.random.random((nb, d)).astype('float32')   # index向量库的向量
xq = np.random.random((nq, d)).astype('float32')   # 待检索的query向量

• step2: 构建index，并将向量添加到index中。这里我们选用暴力检索的方法FlatL2，L2代表构建的index采用的相似度度量方法为L2范数，即欧氏距离

import faiss          
index = faiss.IndexFlatL2(d)    # 创建索引时必须指定向量的维度d
print(index.is_trained)         # 输出为True，代表该类index不需要训练，只需要add向量进去即可
index.add(xb)                   # 将向量库中的向量加入到index中
print(index.ntotal)             # 输出index中包含的向量总数，为100000 

• step3: 检索TopK相似query

k = 4                     # topK的K值
D, I = index.search(xq, k)# xq为待检索向量矩阵，返回的I为每个待检索query最相似TopK的索引list，D为其对应的距离
print(I[:5])
print(D[:5])

得到如下输出，第一个矩阵是索引，第一列0，1，2，3，4表示待检索矩阵自己的索引，从第二距离矩阵也可以看出，自己和自己的距离为0

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

tips 1
在实际使用时，构建索引一般都使用faiss.index_factory 方法，基本所有的index都支持这种构建索引方法，上面步骤2中构建索引可以变成如下方式：

dim, measure = 64, faiss.METRIC_L2
param = 'Flat'
index = faiss.index_factory(dim, param, measure)

dim: 指定向量的维度
param: 是传入index的参数，代表需要构建什么类型的索引
measure：度量方法，目前支持两种，欧氏距离和inner product，即内积。因此，要计算余弦相似度，只需要将向量归一化后，使用内积度量即可。参数为faiss.METRIC_INNER_PRODUCT

tips 2
一些索引可以保存整型的ID，每个向量可以指定一个ID，当查询相似向量时，会返回相似向量的ID及相似度(或距离)。如果不指定，将按照添加的顺序从0开始累加。其中IndexFlatL2不支持指定ID。IndexFlatL2不支持指定id，但是可以通过IDMAP的方式实现，如下代码

ids = [2,10, 100,...]
ids = np.array(ids)
index = faiss.index_factory(768, "IDMap, Flat")
index.add_with_ids(save_embedding, ids)        # 指定id，save_embedding为向量库

下面方法与上面类似

index = faiss.IndexFlatL2(d)
ids = np.arange(100000, 200000)  
index2 = faiss.IndexIDMap(index)
index2.add_with_ids(xb, ids)

4. Faiss常用index优缺点及使用场景

4.1 Flat ：暴力检索

• 优点：该方法是Faiss所有index中最准确的，召回率最高的方法，没有之一；
• 缺点：速度慢，占内存大
• 使用情况：向量候选集很少，在50万以内，并且内存不紧张
  构建方法：

dim, measure = 64, faiss.METRIC_L2
param =  'Flat'
index = faiss.index_factory(dim, param, measure)
index.is_trained                                   # 输出为True
index.add(xb)                                      # 向index中添加向量

4.2 IVFx Flat ：倒排暴力检索

• 优点：该方法是利用倒排技术对暴力检索加速，速度快于Flat
• 缺点：还不是很快，并且检索召回率下降
• 使用情况：同Flat，
• 参数：IVFx中的x是k-means聚类中心的个数
  构建方法：

dim, measure = 64, faiss.METRIC_L2 
param =  'IVF100, Flat'                          # 代表k-means聚类中心为100,   
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)             

4.3 HNSWx

• 优点：该方法为基于图检索的改进方法，检索速度极快，10亿级别秒出检索结果，而且召回率几乎可以媲美Flat，能达到惊人的97%。检索的时间复杂度为loglogn，几乎可以无视候选向量的量级了。并且支持分批导入，极其适合线上任务，毫秒级别体验。
• 缺点：构建索引极慢，占用内存极大（是Faiss中最大的，大于原向量占用的内存大小）
• 参数：HNSWx中的x为构建图时每个点最多连接多少个节点，x越大，构图越复杂，查询越精确，当然构建index时间也就越慢，x取4~64中的任何一个整数。
• 使用情况：不在乎内存，并且有充裕的时间来构建index

构建方法：

dim, measure = 64, faiss.METRIC_L2   
param =  'HNSW64' 
index = faiss.index_factory(dim, param, measure)  
print(index.is_trained)                          # 此时输出为True 
index.add(xb)

4.4 PQx ：乘积量化

• 优点：利用乘积量化的方法，改进了普通k-means，将一个向量的维度切成x段，每段分别进行k-means。因此速度很快，而且占用内存较小，召回率也相对较高
• 缺点：召回率相较于暴力检索，下降较多。
• 使用情况：内存及其稀缺，并且需要较快的检索速度，不那么在意召回率
• 参数：PQx中的x为将向量切分的段数，因此，x需要能被向量维度整除，且x越大，切分越细致，时间复杂度越高

构建方法：

dim, measure = 64, faiss.METRIC_L2 
param =  'PQ16' 
index = faiss.index_factory(dim, param, measure)
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means，
index.train(xb)                                  # 因此需要先训练index，再add向量
index.add(xb)   

4.5 IVFxPQy 倒排乘积量化

• 优点：工业界大量使用此方法，各项指标都均可以接受。
• 缺点：集百家之长，自然也集百家之短
• 使用情况：同PQx
• 参数：IVFxPQy，其中的x和y同上
  构建方法：

dim, measure = 64, faiss.METRIC_L2  
param =  'IVF100, PQ16'
index = faiss.index_factory(dim, param, measure) 
print(index.is_trained)                          # 此时输出为False，因为倒排索引需要训练k-means， 
index.train(xb)                                  # 因此需要先训练index，再add向量 index.add(xb)       

更多参考文献[3]

5. 使用GPU

参考Running on GPUs

5.1 使用单个gpu

res = faiss.StandardGpuResources()  # 声明gpu资源
# 构建一个 flat (CPU) 索引
index_flat = faiss.IndexFlatL2(d)
# 将cpu索引加入到gpu中
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
# 接下来的操作与一般情况类似
gpu_index_flat.add(xb)         # add vectors to the index
print(gpu_index_flat.ntotal)
k = 4                          # we want to see 4 nearest neighbors
D, I = gpu_index_flat.search(xq, k)  # actual search
print(I[:5])                   # neighbors of the 5 first queries
print(I[-5:])                  # neighbors of the 5 last queries

5.2 使用多个GPU

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

cpu_index = faiss.IndexFlatL2(d)
gpu_index = faiss.index_cpu_to_all_gpus(cpu_index) # build the index

gpu_index.add(xb)              # add vectors to the index
print(gpu_index.ntotal)

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

参考

faiss简介及示例
Faiss流程与原理分析
Faiss入门及应用经验记录