简单测试下faiss 检索速度

embedded/2024/11/29 17:58:24/

在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 = dimself.nlist = nlist                      #聚类中心的个数#self.index = faiss.IndexFlatL2(dim)    # build the indexquantizer = 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 searchself.index = faiss.IndexIVFFlat(quantizer, dim, self.nlist, faiss.METRIC_L2)try:if gpu>=0:if gpu==0:# use a single GPUres = 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_indexexcept :pass# data self.xb = None# 返回当前总共有多少个值def curr_items ():# self.index.ntotalreturn self.xb.shape[0]# 清空数据def reset (self):passself.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 searchreturn 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 += vreturn 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 秒


http://www.ppmy.cn/embedded/141535.html

相关文章

如何使用PHP爬虫获取店铺详情：一篇详尽指南

如何使用PHP爬虫获取店铺详情：一篇详尽指南

在数字化时代，数据的价值不言而喻。对于企业来说，获取竞争对手的店铺详情、顾客评价等信息对于市场分析和决策至关重要。PHP作为一种广泛使用的服务器端脚本语言，结合其强大的库支持，使得编写爬虫程序变得简单而高效。本文将详细介…
阅读更多...
JVM_栈详解一

JVM_栈详解一

1、栈的存储单位 **栈中存储什么？**， 每个线程都有自己的栈，栈中的数据都是以栈帧（Stack Frame）的格式存在。在这个线程上正在执行的每个方法都各自对应一个栈帧（Stack Frame）。 栈帧是一个内存…
阅读更多...
38 基于单片机的宠物喂食(ESP8266、红外、电机)

38 基于单片机的宠物喂食(ESP8266、红外、电机)

目录 一、主要功能 二、硬件资源 三、程序编程 四、实现现象 一、主要功能 基于STC89C52单片机，采用L298N驱动连接P2.3和P2.4口进行电机驱动， 然后串口连接P3.0和P3.1模拟ESP8266， 红外传感器连接ADC0832数模转换器连接单片机的P1.0~P1.…
阅读更多...
基于SpringBoot实现的民宿管理系统(代码+论文)

基于SpringBoot实现的民宿管理系统(代码+论文)

🎉博主介绍：Java领域优质创作者，阿里云博客专家，计算机毕设实战导师。专注Java项目实战、毕设定制/协助 📢主要服务内容：选题定题、开题报告、任务书、程序开发、项目定制、论文辅导 💖精彩专栏…
阅读更多...
【数据库系列】Flyway详解及详细使用步骤

【数据库系列】Flyway详解及详细使用步骤

什么是Flyway？ Flyway是一个开源的数据库迁移工具，旨在帮助开发者管理数据库版本和迁移。它支持多种数据库，包括MySQL、PostgreSQL、Oracle和SQL Server等。Flyway通过版本控制的方式，确保数据库的结构和数据与代码库中的版本保持…
阅读更多...
AI生成一个Supermap GIS开发大赛的一个作品

AI生成一个Supermap GIS开发大赛的一个作品

2024年Supermap GIS大赛！加油！ 参赛作品设计应充分展示SuperMap系列产品在地理信息系统（GIS）领域的强大功能和广泛应用。以下是一个基于SuperMap软件系列设计的参赛作品概述，旨在体现其数据处理、分析、制图及发布等核…
阅读更多...
VUE 生成 二维码（qrcodejs2-fix），条形码（jsbarcode）

VUE 生成 二维码（qrcodejs2-fix），条形码（jsbarcode）

二维码 需要用到依赖&#xff1a;qrcodejs2-fix 安装依赖 npm i qrcodejs2-fix 代码部分 <template><div><div id"codes" ref"codes"></div></div> </template><script setup> import { ref, onMounted } …
阅读更多...
django实现paypal订阅记录

django实现paypal订阅记录

开发者链接 登录开发者 沙箱账号链接 沙箱账号链接 1. 创建沙箱应用 2. 记录 Client ID和 Secret 后面有用 然后点击进去创建回调url和回调事件&#xff0c;可以全选事件也可以选择你需要的 3.拿沙箱账号信息去登录 4.登录后创建订阅产品内容。记住产品id 5.配置django后端信…
阅读更多...
最新文章

Copyright @ 2022~2023