在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 秒