简单测试下faiss 检索速度

在NLP的应用中，经常需要用到对向量的搜索，如果向量的数量级非常大，比如1千万，甚至上亿条，普通的方式就满足不了生产需要了，falcebook开源的faiss框架能够解决"海量向量搜索"的问题。faiss是为稠密向量提供高效相似度搜索和聚类的框架。由Facebook AI Research研发。 具有以下特性。

• 1、提供多种检索方法
• 2、速度快
• 3、可存在内存和磁盘中
• 4、C++实现，提供Python封装调用。
• 5、大部分算法支持GPU实现

GitHub - xmxoxo/faiss_test: faiss packaging and faiss speed test

参考了这个项目里做一个对比实验，来看一下faiss的使用及速度。 随机生成10万个768维的向量，加载到内存中， 然后分别用普通的暴力搜索和faiss搜索两种方式去搜索，对比搜索的平时用时。

按照相同规则生成随机向量，一般句子向量维度是768

# 随机生成向量 1百万   # total = 1000000
# 向量的维度   # dim = 1024 #768
print('随机生成%d个向量，维度：%d' % (total, dim), flush=True)
#rng = np.random.RandomState(0)
#X = rng.random_sample((total, dim))
X = np.random.random((total, dim))

以下是对每一部分的详细分析：

创建 VecSearch 实例

vs = VecSearch(dim=dim, gpu=gpu)

• VecSearch : 类的实例化，dim 是向量的维度，gpu 指定是否使用GPU。

• 这一步会初始化FAISS索引并配置GPU资源（如果指定了）。

  class VecSearch:
  def init(self):
  self.dicts = {}

    # 返回当前总共有多少个值
    def curr_items ():
        return len(self.dicts)
  
    # 添加文档
    def add_doc (self, key, vector):
        self.dicts[key] = vector
  
    # 查找向量,
    # 返回结果为 距离[D], 索引[I]
    def search(self, query, top=5):
        # 返回结果，结构为：[sim, key]
        ret = np.zeros((top,2))
        # 计算余弦相似度最大值
        for key, value in self.dicts.items():
            sim = CosSim_dot(query, value)
            #sim = CosSim(query, value)
            #sim = CosSim_sk(query, value)
            #sim = cosine(query, value)
            #print(sim)
            if sim > ret[top-1][0]:
                b = np.array([[sim, key]]).astype('float32')
                ret = np.insert(ret, 0, values=b, axis=0)
                # 重新排序后截取
                idex = np.lexsort([-1*ret[:,0]])
                ret = ret[idex, :]
                ret = ret[:top,]
                #print(ret)
                #print('-'*40)
        return ret[:,0], ret[:,1].astype('int')
  

上面**search 方法** : 用于查找与给定查询向量 query 最相似的文档向量，返回前 top 个最相似的结果。ret 为一个形状为 (top, 2) 的数组，用于存储相似度和对应的文档索引。

遍历存储的向量 : 对字典中的每个文档向量，计算与查询向量的余弦相似度 sim。可以选择不同的相似度计算方法（如 CosSim, cosine 等），但此处使用的是 CosSim_dot。

更新相似度结果 : 如果当前相似度 sim 大于 ret 中最小的相似度，则将其插入到 ret 的开头。

排序和截取 : 使用 np.lexsort 按相似度对 ret 进行排序，并截取前 top 个结果。

ret[:, 0], ret[:, 1].astype('int')

返回相似度和对应索引: 方法返回两个数组，分别是相似度和文档的索引（转换为整数）。

添加向量

ret = vs.add(X)

• add(X) : 将生成的向量 X 添加到 VecSearch 实例中。
• X 是一个形状为 (N, dim) 的数组，其中 N 是向量的数量。
• 返回值 ret 是一个元组，包含添加前后的索引范围，例如 (起始索引, 结束索引)。

训练和索引

vs.reindex()

• reindex(): 训练FAISS索引并将数据添加到索引中。
• 此步骤是必要的，因为在向量添加后，FAISS需要训练索引以便能够有效执行后续的搜索操作。

计算创建时间

end = time.time()
total_time = end - start
print('创建用时:%4f秒' % total_time)

• 这里通过记录结束时间 end 和开始时间 start 的差值来计算创建索引和添加向量所花费的时间。
• 使用 print 输出创建索引的总时间，格式化为小数点后四位。

查看内存使用情况

import os, psutil
process = psutil.Process(os.getpid())
print('Used Memory:', process.memory_info().rss / 1024 / 1024, 'MB')

• psutil: 用于获取系统和进程信息的库。
• os.getpid() 获取当前进程的ID。
• process.memory_info().rss 返回进程当前使用的物理内存（RSS: Resident Set Size）。
• 将内存使用量从字节转换为MB，并打印出来。

获取当前进程的内存使用情况

process = psutil.Process(os.getpid())
print('Used Memory:', process.memory_info().rss / 1024 / 1024, 'MB')

• 这段代码使用 psutil 库来获取当前运行进程的内存使用情况。
• os.getpid() 获取当前Python程序的进程ID。
• process.memory_info().rss 获取该进程的常驻集大小（RSS），即当前使用的物理内存量（以字节为单位）。
• 将字节转换为MB（通过除以1024两次），并打印出来，便于查看内存占用情况。

单条查询测试的开始

print('单条查询测试'.center(40,'-'))

• 打印一行文本，中心对齐并用 - 符号填充，便于在输出中分隔不同的测试部分。

生成查询向量

Q = np.random.random((test_times, dim))
Q[:, 0] += np.arange(test_times) / test_times

• 生成一个形状为 (test_times, dim) 的随机数组 Q，其中 test_times 是测试的次数，dim 是向量的维度。
• 通过 Q[:, 0] += np.arange(test_times) / test_times，对第一列进行线性调整，使得查询向量 Q 的第一维数据呈现一定的趋势。这有助于在搜索时产生更具代表性的查询结果。

执行单条查询

q = Q[0]
start = time.time()
D, I = vs.search(q, top=top_n, nprobe=10)

• q = Q[0]: 选择生成的查询向量中的第一条作为单条查询。
• start = time.time(): 记录查询开始的时间，用于后续计算查询所花费的时间。
• D, I = vs.search(q, top=top_n, nprobe=10) : 使用 VecSearch 实例 vs 执行搜索：
  • q 是要查询的向量。
  • top=top_n 指定返回的最相似向量的数量。
  • nprobe=10 指定在查询时要探测的聚类中心的数量，这会影响查询的准确性和速度。

测试下暴力查询 的检索速度

[root@node126 embeding]# /opt/miniconda3/envs/rag/bin/python vector_search_force.py
===========大批量向量余弦相似度计算-[暴力版]===========
随机生成100000个向量，维度：768
正在创建搜索器...
添加用时:0.034196秒
Used Memory: 682.5859375 MB
-----------------单条查询测试-----------------
搜索结果: [0.78086126 0.77952558 0.77949381 0.77775592 0.77547914] [ 541  443 1472  370  209]
显示查询结果，并验证余弦相似度...
索引号:  541, 距离:0.780861
索引号:  443, 距离:0.779526
索引号: 1472, 距离:0.779494
索引号:  370, 距离:0.777756
索引号:  209, 距离:0.775479
-----------------批量查询测试-----------------
批量测试次数：100 次，请稍候...
总用时:59 秒, 平均用时:590.072079 毫秒

分析下代码项目faiss检索代码

使用faiss 实现VecSearch

class VecSearch:
    def __init__(self, dim=10, nlist=100, gpu=-1):
        self.dim = dim
        self.nlist = nlist                      #聚类中心的个数
        #self.index = faiss.IndexFlatL2(dim)    # build the index
        quantizer = faiss.IndexFlatL2(dim)      # the other index

        # faiss.METRIC_L2: faiss定义了两种衡量相似度的方法(metrics)，
        # 分别为faiss.METRIC_L2 欧式距离、 faiss.METRIC_INNER_PRODUCT 向量内积
        # here we specify METRIC_L2, by default it performs inner-product search
        self.index = faiss.IndexIVFFlat(quantizer, dim, self.nlist, faiss.METRIC_L2)
        
        try:
            if gpu>=0:
                if gpu==0:
                    # use a single GPU
                    res = faiss.StandardGpuResources()  
                    gpu_index = faiss.index_cpu_to_gpu(res, 0, self.index)
                else:
                    gpu_index = faiss.index_cpu_to_all_gpus(self.index)
            
                self.index = gpu_index
        except :
            pass
        
        # data 
        self.xb = None
    
    # 返回当前总共有多少个值
    def curr_items ():
        # self.index.ntotal
        return self.xb.shape[0]

    # 清空数据
    def reset (self):
        pass
        self.xb = None
        

    # 添加向量，可批量添加，编号是按添加的顺序；
    # 参数: vector, 大小是(N, dim)
    # 返回结果：索引号区间, 例如 (0,8), (20,100)
    def add (self, vector):
        if not vector.dtype == 'float32':
            vector = vector.astype('float32')
        
        if self.xb is None:
            prepos = 0
            # vector = vector[np.newaxis, :]   
            self.xb = vector.copy()
        else:
            prepos = self.xb.shape[0]
            self.xb = np.vstack((self.xb,q))
        
        return (prepos, self.xb.shape[0]-1)

    # 添加后开始训练
    def reindex(self):
        self.index.train(self.xb)
        self.index.add(self.xb)                  # add may be a bit slower as well
    

    # 查找向量, 可以批量查找，
    # 参数：query (N,dim)
    # 返回： 距离D,索引号I  两个矩阵
    def search(self, query, top=5, nprobe=1):
        # 查找聚类中心的个数，默认为1个。
        self.index.nprobe = nprobe #self.nlist 

        # 如果是单条查询，把向量处理成二维 
        #print(query.shape)
        if len(query.shape)==1:
            query = query[np.newaxis, :]
        #print(query.shape)
        # 查询
        if not query.dtype == 'float32':
            query = query.astype('float32')
        D, I = self.index.search(query, top)     # actual search
        return D, I

• FAISS索引:

  • 使用 faiss.IndexFlatL2 创建量化器。L2距离用于计算向量之间的距离。
  • faiss.IndexIVFFlat 创建用于高效检索的IVF索引。

向量搜索 search

def search(self, query, top=5, nprobe=1):
    self.index.nprobe = nprobe  # 设置要探测的聚类中心数量

    if len(query.shape) == 1:
        query = query[np.newaxis, :]  # 确保查询是二维的
    
    if not query.dtype == 'float32':
        query = query.astype('float32')
    
    D, I = self.index.search(query, top)  # 执行搜索
    return D, I

• 功能: 执行向量搜索，返回最相似的向量。

• 参数:

  • query: 要搜索的向量。
  • top: 返回最相似的向量数量（默认5）。
  • nprobe: 要探测的聚类中心数量（默认1）。

• 过程:

  • 设置 self.index.nprobe 用于搜索时的聚类中心数量。
  • 检查查询向量的维度，确保其为二维（batch size, dim）。
  • 确保查询向量为 float32 类型。
  • 调用 self.index.search(query, top) 执行搜索，返回距离 D 和索引 I。

代码中使用到的几个方法

1. seg_vector 函数

def seg_vector(txt, dict_vector, emb_size=768):
    seg_v = np.zeros(emb_size)
    for w in txt:
        if w in dict_vector.keys():
            v = dict_vector[w]
            seg_v += v
    return seg_v

• 功能 : 将输入的文本 txt 转换为一个句向量。这个句向量是通过对文本中单词的向量进行简单相加得到的。

• 参数:

  • txt: 输入的文本，可以是一个单词的列表或字符串。
  • dict_vector: 一个字典，映射单词到其对应的向量（通常是预训练的词向量）。
  • emb_size: 向量的维度，默认为768。

• 过程:

  • seg_v = np.zeros(emb_size): 创建一个全零的向量，长度为 emb_size，用于存储句向量。
  • for w in txt: 遍历文本中的每个单词 w。
  • if w in dict_vector.keys(): 检查单词 w 是否在字典中。
  • v = dict_vector[w]: 如果在，获取该单词对应的向量 v。
  • seg_v += v: 将单词向量 v 加到句向量 seg_v 上。

• 返回 : 最终返回的 seg_v 是文本 txt 的句向量。

2. CosSim 函数

def CosSim(a, b):
    return 1 - cosine(a, b)

• 功能 : 计算两个向量 a 和 b 之间的余弦相似度。

• 过程:

  • 使用 scipy.spatial.distance 中的 cosine 函数来计算余弦距离（即1减去余弦相似度）。
  • 余弦相似度的值范围在[-1, 1]之间，通常在实际应用中我们更关注0到1之间的值，所以通过 1 - cosine(a, b) 转换为相似度。

• 返回 : 返回 a 和 b 之间的余弦相似度，值越接近1，表示相似度越高。

3. CosSim_dot 函数

def CosSim_dot(a, b):
    score = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
    return score

• 功能 : 计算两个向量 a 和 b 的余弦相似度，使用内积和范数。

• 过程:

  • np.dot(a, b): 计算向量 a 和 b 的内积。
  • np.linalg.norm(a): 计算向量 a 的范数（即长度）。
  • np.linalg.norm(b): 计算向量 b 的范数。
  • score = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)): 使用内积和向量的范数计算余弦相似度。

• 返回: 返回计算得到的余弦相似度，值范围也是[-1, 1]。

测试下faiss 的检索速度

[root@node126 embeding]# /opt/miniconda3/envs/rag/bin/python vector_search_faiss.py
=========大批量向量余弦相似度计算-[faiss版]==========
随机生成100000个向量，维度：768
正在创建搜索器...
GPU使用情况:不使用
创建用时:5.961992秒
Used Memory: 1613.765625 MB
-----------------单条查询测试-----------------
显示查询结果，并验证余弦相似度...
索引号: 1058, 距离:114.878304, 余弦相似度:0.771127
索引号:  541, 距离:115.051338, 余弦相似度:0.780861
索引号:  370, 距离:115.715919, 余弦相似度:0.777756
索引号:  209, 距离:115.731323, 余弦相似度:0.775479
索引号: 1472, 距离:115.832909, 余弦相似度:0.779494
总用时:10毫秒
-----------------批量查询测试-----------------
正在批量测试10000次，每次返回Top 5，请稍候...
总用时:13576毫秒, 平均用时:1.357649毫秒
----------------------------------------

看到搜索速度 13576毫秒 远远优于 59 秒