faiss向量数据库测《三体》全集,这家国产AI加速卡把性能提了7倍

科技有夕小瑶 2024-08-20 14:44:57

在人工智能和机器学习技术的飞速发展中,向量数据库在处理高维数据方面扮演着日益重要的角色。近年来,随着大型模型的流行,向量数据库技术也得到了进一步的发展和完善。

向量数据库为大型模型提供了一个高效的数据管理和检索平台,使得这些模型能够更加高效地处理非结构化数据,并在各种应用场景中发挥其潜力,执行复杂的查询和分析任务。

Faiss(Facebook AI Similarity Search)是一个用于高效相似性搜索和密集向量聚类的库,它广泛应用于图像检索、推荐系统和自然语言处理等领域。然而,随着数据规模的不断扩大和维度的增加,如何在保证搜索质量的同时提升搜索速度成为了一个挑战。海光DCU(Data Center Unit)作为一种高性能的计算加速解决方案,能够有效提升Faiss向量数据库的搜索性能。

Faiss向量数据库简介

Faiss是由Facebook AI Research团队开发的一个开源库,专门用于高效地进行大规模向量的相似性搜索和聚类。它支持对十亿级别的向量进行搜索,是目前较为成熟的近似近邻搜索库之一。Faiss用C++编写,并提供了与Numpy紧密结合的Python接口,不仅支持CPU计算,对一些核心算法还支持GPU计算。

海光DCU简介

海光DCU(Data Center Unit)是一款高效通用的GPGPU加速卡,专为人工智能和科学计算任务设计。它在兼容性、软件生态和市场应用方面展现出卓越的价值。海光DCU全面兼容“类 CUDA”环境。这种强大的兼容性为用户提供了在AI和大数据处理领域的强大计算服务能力,其在国产加速卡领域中的生态兼容性处于领先地位。

DCU环境部署

本次测试使用了一台装备有两张海光Z100L加速卡的服务器X7840H0,服务器操作系统为Ubuntu 22.04.1 LTS。

准备开发测试环境,相关的程序和文档可以通过光和开发者社区获取,地址是https://developer.hpccube.com。

在服务器系统上部署开发测试环境,用户可以通过点击页面上的资源工具访问驱动、DTK、DAS、镜像等资源的下载界面。

《开发环境安装部署手册》可以通过点击DTK Toolkit下载地址,然后选择最新的latest,然后选择Document目录获取。除了开发环境安装部署手册外,还有开发环境使用手册、兼容性手册等常用的说明文档。

《开发环境安装部署手册》中包含了多个常用系统下的基础环境部署,可以根据使用的系统选择对应的环境部署流程:

根据测试机服务器的操作系统版本,本次测试选择Ubuntu20.04.1操作系统基础环境部署。

按照手册中要求的首先安装驱动以及DTK的依赖包,然后安装驱动程序和DTK,设备的DCU开发测试环境即可部署完成。环境部署完成后输入hy-smi指令即可查询DCU的使用信息:

除了使用物理机的系统环境开发测试外,还可以使用官方提供的基础环境镜像,镜像下载地址:https://sourcefind.cn/#/main-page。

使用官方提供的镜像可以节省大量基础环境的部署工作。本次测试就使用到了名称为1.13.1-centos7.6-dtk-23.04.1-py38-latest的镜像,镜像内已部署好了pytorch等相关的第三方包。然后安装光和开发者社区中提供的faiss安装包以及测试代码所需的pandas等三方包就可以进入下一步准备faiss的测试程序。

搜索性能测试

为了测试faiss的搜索效率,本次测试以文本相似度搜索为例,分别在CPU和GPU场景下进行测试。测试流程包括将批量文本数据导入faiss向量数据库,然后搜索一段文本中不存在的数据,并取多次测试的平均值进行对比。

将文本数据转换为向量数据需要用到Embedding嵌入模型,本次测试中选择了效果出色的shibing624/text2vec-base-chinese。

文本内容本次测试选择了《三体》全集,文本存储在三体.txt文件中。由于Embedding嵌入模型的输入长度限制,首先需要将文本内容进行分段再传入嵌入模型。然后将嵌入模型转换完成的向量数据使用numpy存储在本地data.npy文件中,用来方便后续测试。代码如下:

import warningswarnings.simplefilter(action='ignore', category=FutureWarning)import pandas as pddf=pd.read_csv("三体.txt",encoding='utf-8',sep="#",header=None, names=["sentence"])print(df)from sentence_transformers import SentenceTransformermodel=SentenceTransformer('shibing624/text2vec-base-chinese')sentences =df['sentence'].tolist()sentence_embeddings=model.encode(sentences)print("数据向最维度:")print(sentence_embeddings.shape)save_file = "data.npy"import numpy as npnp.save(save_file,sentence_embeddings)import osfile_size = os.path.getsize(save_file)print("保存数据文件:%7.3f MB"%(file_size/1024/1024))

运行代码之后打印信息如下:

向量数据准备好之后使用faiss分别加载三体全集和data.npy向量数据,然后使用faiss中提供的IndexFlatL2索引方式加载这些向量数据,然后在搜索“大史喜欢抽烟”这几个原文中没有的文本。faiss_test.py测试代码如下:

import faissimport numpy as npimport pandas as pdimport warningswarnings.simplefilter(action='ignore', category=FutureWarning)print("load 三体.txt...")df = pd.read_csv("三体.txt", encoding='utf-8', sep="#", header=None, names=["sentence"])print("load vector data...")sentence_embeddings = np.load("data.npy")dimension = sentence_embeddings.shape[1]index = faiss.IndexFlatL2(dimension)index.add(sentence_embeddings)import timefrom sentence_transformers import SentenceTransformermodel = SentenceTransformer('shibing624/text2vec-base-chinese')topk = 5words = ["大史喜欢抽烟"]search = model.encode(words)print("search: " + str(words))costs = []for i in range(10): to = time.time() D, I = index.search(search, topk) ti = time.time() costs.append(ti - to)print(D)print(I)print(df['sentence'].iloc[I[0]])print("平均耗时 %7.3f ms" % ((sum(costs) / len(costs)) * 1000.0))

使用GPU的方式搜索可以将上面代码中的index使用index_cpu_to_all_gpus的方法将索引数据创建在GPU中,然后构建索引数据。faiss_gpu_test.py代码如下:

import faiss time warningsimport numpy as npimport pandas as pdwarnings.simplefilter(action='ignore', category=FutureWarning)print("load 三体.txt...")df = pd.read_csv("三体.txt", encoding='utf-8', sep="#", header=None, names=["sentence"])print("load vector data...")sentence_embeddings = np.load("data.npy")dimension = sentence_embeddings.shape[1]index = faiss.IndexFlatL2(dimension)ngpus = faiss.get_num_gpus()print("number of GPU:", ngpus)gpu_index = faiss.index_cpu_to_all_gpus(index)gpu_index.add(sentence_embeddings)from sentence_transformers import SentenceTransformermodel = SentenceTransformer('shibing624/text2vec-base-chinese')topk = 5words = ["大史喜欢抽烟"]search = model.encode(words)print("search: " + str(words))costs = []for i in range(10): to = time.time() D, I = gpu_index.search(search, topk) ti = time.time() costs.append(ti - to)print(D)print(I)print(df['sentence'].iloc[I[0]])print("平均耗时 %7.3f ms" % ((sum(costs) / len(costs)) * 1000.0))

在服务器环境中分别运行faiss_test.py和faiss_gpu_test.py即可获取到faiss的搜索结果:

分析汇总

对两份代码的运行结果进行对比可以明显看到海光DCU的加速效果明显,较CPU索引的方式提高了7倍左右的性能。

本次测试使用到的文本数据量较低,随着数据量的增加,DCU的加速效果会更加明显。测试代码中的索引方式使用到了faiss中最基本的IndexFlatL2,它使用 L2 距离(欧氏距离)进行暴力搜索(brute-force search),适用于向量数量较小的情况。由于它在内存中存储所有向量,因此当向量数量较大时,内存开销会很大。除此之外faiss中常用的还有IndexIVFFlat、IndexIVFPQ等索引方式可以显著减少索引的内存资源占用。

0 阅读:1

科技有夕小瑶

简介:感谢大家的关注