在人工智能和机器学习技术的飞速发展中,向量数据库在处理高维数据方面扮演着日益重要的角色。近年来,随着大型模型的流行,向量数据库技术也得到了进一步的发展和完善。
向量数据库为大型模型提供了一个高效的数据管理和检索平台,使得这些模型能够更加高效地处理非结构化数据,并在各种应用场景中发挥其潜力,执行复杂的查询和分析任务。
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 warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
import pandas as pd
df=pd.read_csv("三体.txt",encoding='utf-8',sep="#",header=None, names=["sentence"])
print(df)
from sentence_transformers import SentenceTransformer
model=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 np
np.save(save_file,sentence_embeddings)
import os
file_size = os.path.getsize(save_file)
print("保存数据文件:%7.3f MB"%(file_size/1024/1024))
运行代码之后打印信息如下:
向量数据准备好之后使用faiss分别加载三体全集和data.npy向量数据,然后使用faiss中提供的IndexFlatL2索引方式加载这些向量数据,然后在搜索“大史喜欢抽烟”这几个原文中没有的文本。faiss_test.py测试代码如下:
import faiss
import numpy as np
import pandas as pd
import warnings
warnings.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 time
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('shibing624/text2vec-base-chinese')
topk = 5
words = ["大史喜欢抽烟"]
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 warnings
import numpy as np
import pandas as pd
warnings.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 SentenceTransformer
model = SentenceTransformer('shibing624/text2vec-base-chinese')
topk = 5
words = ["大史喜欢抽烟"]
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等索引方式可以显著减少索引的内存资源占用。