cuvs

cuvs

高性能GPU向量搜索与聚类框架

cuVS是一个开源的GPU向量搜索和聚类框架,专注于提供高性能的近似最近邻搜索和聚类功能。它支持C、C++、Python和Rust等多种编程语言,实现了CAGRA等先进算法。开发者可以直接使用cuVS,也可将其集成到其他系统中,从而在向量相似度搜索和聚类任务中充分利用GPU加速能力。

cuVSGPU向量搜索聚类算法RAPIDSGithub开源项目

<div align="left"><img src="https://yellow-cdn.veclightyear.com/ab5030c0/5e553fe1-0574-4870-9f9c-f8fea5811283.png" width="90px"/> cuVS: GPU上的向量搜索和聚类</div>

[!note] cuVS是一个新的库,主要源于RAPIDS RAFT数据挖掘原语库中的近似最近邻和聚类算法。RAPIDS RAFT目前包含了cuVS中最全面的近似最近邻和聚类算法版本。我们正在将算法从RAFT迁移到cuVS,但如果您不确定使用哪个,请考虑以下几点:

  1. RAFT为所有近似最近邻和聚类算法提供C++和Python API。
  2. cuVS对不同语言提供越来越多的支持,包括C、C++、Python和Rust。我们将在未来为cuVS添加更多语言支持,但不会改进RAFT的语言支持。
  3. 一旦RAFT的所有近似最近邻和聚类算法都迁移到cuVS,RAFT的API将被弃用并最终完全移除。移除后,RAFT将成为一个轻量级的仅头文件库。在此期间,如果不需要额外的语言支持,使用RAFT也没有问题。

目录

  1. 有用资源
  2. 什么是cuVS?
  3. 安装cuVS
  4. 开始使用
  5. 贡献
  6. 参考文献

有用资源

什么是cuVS?

cuVS包含了几种在GPU上运行近似最近邻和聚类的最先进算法实现。它可以直接使用,也可以通过集成了它的各种数据库和其他库使用。cuVS的主要目标是简化GPU在向量相似度搜索和聚类中的使用。

安装cuVS

cuVS提供了可以通过conda安装的预构建包。cuVS支持的不同语言有不同的包:

PythonC/C++
cuvslibcuvs, libcuvs-static

稳定版本

建议使用mamba来安装所需的包。以下命令将安装Python包。您可以用上表中的任何包替换cuvs

mamba install -c conda-forge -c nvidia -c rapidsai cuvs

每日构建版本

如果安装尚未发布的版本,可以将rapidsai通道替换为rapidsai-nightly

mamba install -c conda-forge -c nvidia -c rapidsai-nightly cuvs=24.10

有关安装cuVS和从源代码构建的更多信息,请参阅构建和安装指南

开始使用

以下代码片段为CAGRA算法训练一个近似最近邻索引。

Python API

from cuvs.neighbors import cagra dataset = load_data() index_params = cagra.IndexParams() index = cagra.build(build_params, dataset)

C++ API

#include <cuvs/neighbors/cagra.hpp> using namespace cuvs::neighbors; raft::device_matrix_view<float> dataset = load_dataset(); raft::device_resources res; cagra::index_params index_params; auto index = cagra::build(res, index_params, dataset);

有关C++ API的更多示例,请参阅代码库中的cpp/examples目录。

C API

#include <cuvs/neighbors/cagra.h> cuvsResources_t res; cuvsCagraIndexParams_t index_params; cuvsCagraIndex_t index; DLManagedTensor *dataset; load_dataset(dataset); cuvsResourcesCreate(&res); cuvsCagraIndexParamsCreate(&index_params); cuvsCagraIndexCreate(&index); cuvsCagraBuild(res, index_params, dataset, index); cuvsCagraIndexDestroy(index); cuvsCagraIndexParamsDestroy(index_params); cuvsResourcesDestroy(res);

Rust API

use cuvs::cagra::{Index, IndexParams, SearchParams}; use cuvs::{ManagedTensor, Resources, Result}; use ndarray::s; use ndarray_rand::rand_distr::Uniform; use ndarray_rand::RandomExt; /// 展示如何使用CAGRA索引和搜索数据的示例 fn cagra_example() -> Result<()> { let res = Resources::new()?; // 创建一个新的随机数据集进行索引 let n_datapoints = 65536; let n_features = 512; let dataset = ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0)); // 构建cagra索引 let build_params = IndexParams::new()?; let index = Index::build(&res, &build_params, &dataset)?; println!( "将{}x{}个数据点索引到cagra索引中", n_datapoints, n_features ); // 使用数据集中的前4个点作为查询:将测试我们是否能够将它们作为自己的最近邻找回 let n_queries = 4; let queries = dataset.slice(s![0..n_queries, ..]); let k = 10; // CAGRA搜索API要求查询和输出在设备内存上 // 复制查询数据,并为距离/邻居输出分配新的设备内存 let queries = ManagedTensor::from(&queries).to_device(&res)?; let mut neighbors_host = ndarray::Array::<u32, _>::zeros((n_queries, k)); let neighbors = ManagedTensor::from(&neighbors_host).to_device(&res)?;

let mut distances_host = ndarray::Array::<f32, _>::zeros((n_queries, k)); let distances = ManagedTensor::from(&distances_host).to_device(&res)?;

let search_params = SearchParams::new()?;

index.search(&res, &search_params, &queries, &neighbors, &distances)?;

// 复制回主机内存 distances.to_host(&res, &mut distances_host)?; neighbors.to_host(&res, &mut neighbors_host)?;

// 最近邻应该是它们自己,因为查询来自数据集 println!("邻居 {:?}", neighbors_host); println!("距离 {:?}", distances_host); Ok(()) }


## 贡献

如果你有兴趣为cuVS库做出贡献,请阅读我们的[贡献指南](docs/source/contributing.md)。有关开发者指南、工作流程和原则的详细信息,请参阅[开发者指南](docs/source/developer_guide.md)。

## 参考文献

在引用cuVS时,请考虑引用此Github仓库。
```bibtex
@misc{rapidsai,
  title={Rapidsai/cuVS: GPU上的向量搜索和聚类。},
  url={https://github.com/rapidsai/cuvs},
  journal={GitHub},
  publisher={Nvidia RAPIDS},
  author={Rapidsai},
  year={2024}
}

如果引用CAGRA,请考虑以下bibtex:

@misc{ootomo2023cagra, title={CAGRA:用于GPU的高度并行图构建和近似最近邻搜索}, author={Hiroyuki Ootomo and Akira Naruse and Corey Nolet and Ray Wang and Tamas Feher and Yong Wang}, year={2023}, eprint={2308.15136}, archivePrefix={arXiv}, primaryClass={cs.DS} }

如果引用k-选择例程,请考虑以下bibtex:

@proceedings{10.1145/3581784, title = {SC '23:高性能计算、网络、存储和分析国际会议论文集}, year = {2023}, isbn = {9798400701092}, publisher = {美国计算机协会}, address = {纽约,纽约,美国}, abstract = {始于1988年,SC会议已成为学术界、产业界和政府研究人员和从业者每年分享信息并促进合作以推进高性能计算(HPC)、网络、存储和分析最新技术的重要聚会。}, location = {, 科罗拉多州丹佛, 美国, } }

如果引用最近邻下降API,请考虑以下bibtex:

@inproceedings{10.1145/3459637.3482344, author = {Wang, Hui and Zhao, Wan-Lei and Zeng, Xiangxiang and Yang, Jianye}, title = {通过基于GPU的NN-Descent快速构建K-NN图}, year = {2021}, isbn = {9781450384469}, publisher = {美国计算机协会}, address = {纽约,纽约,美国}, url = {https://doi.org/10.1145/3459637.3482344}, doi = {10.1145/3459637.3482344}, abstract = {NN-Descent是一种经典的k-NN图构建方法。由于其效率和通用性,它在机器学习、计算机视觉和信息检索任务中仍被广泛使用。然而,当前的设计仅在CPU上表现良好。在本文中,NN-Descent被重新设计以适应GPU架构。提出了一种称为选择性更新的新图更新策略。它显著减少了GPU核心和GPU全局内存之间的数据交换,这是GPU计算架构下的处理瓶颈。这种重新设计充分利用了GPU硬件的并行性。同时,NN-Descent的通用性和简单性得到了很好的保留。此外,还提出了一种在GPU上高效合并k-NN图的程序。它使得构建高质量的超出GPU内存数据集的k-NN图变得可行。我们的方法比单线程NN-Descent快100-250倍,比现有的基于GPU的方法快2.5-5倍,这是我们在百万级和十亿级数据集上测试的结果。}, booktitle = {第30届ACM国际信息与知识管理会议论文集}, pages = {1929–1938}, numpages = {10}, keywords = {高维, nn-descent, gpu, k-最近邻图}, location = {虚拟活动,昆士兰,澳大利亚}, series = {CIKM '21} }

编辑推荐精选

商汤小浣熊

商汤小浣熊

最强AI数据分析助手

小浣熊家族Raccoon,您的AI智能助手,致力于通过先进的人工智能技术,为用户提供高效、便捷的智能服务。无论是日常咨询还是专业问题解答,小浣熊都能以快速、准确的响应满足您的需求,让您的生活更加智能便捷。

imini AI

imini AI

像人一样思考的AI智能体

imini 是一款超级AI智能体,能根据人类指令,自主思考、自主完成、并且交付结果的AI智能体。

Keevx

Keevx

AI数字人视频创作平台

Keevx 一款开箱即用的AI数字人视频创作平台,广泛适用于电商广告、企业培训与社媒宣传,让全球企业与个人创作者无需拍摄剪辑,就能快速生成多语言、高质量的专业视频。

即梦AI

即梦AI

一站式AI创作平台

提供 AI 驱动的图片、视频生成及数字人等功能,助力创意创作

扣子-AI办公

扣子-AI办公

AI办公助手,复杂任务高效处理

AI办公助手,复杂任务高效处理。办公效率低?扣子空间AI助手支持播客生成、PPT制作、网页开发及报告写作,覆盖科研、商业、舆情等领域的专家Agent 7x24小时响应,生活工作无缝切换,提升50%效率!

TRAE编程

TRAE编程

AI辅助编程,代码自动修复

Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。

AI工具TraeAI IDE协作生产力转型热门
蛙蛙写作

蛙蛙写作

AI小说写作助手,一站式润色、改写、扩写

蛙蛙写作—国内先进的AI写作平台,涵盖小说、学术、社交媒体等多场景。提供续写、改写、润色等功能,助力创作者高效优化写作流程。界面简洁,功能全面,适合各类写作者提升内容品质和工作效率。

AI辅助写作AI工具蛙蛙写作AI写作工具学术助手办公助手营销助手AI助手
问小白

问小白

全能AI智能助手,随时解答生活与工作的多样问题

问小白,由元石科技研发的AI智能助手,快速准确地解答各种生活和工作问题,包括但不限于搜索、规划和社交互动,帮助用户在日常生活中提高效率,轻松管理个人事务。

热门AI助手AI对话AI工具聊天机器人
Transly

Transly

实时语音翻译/同声传译工具

Transly是一个多场景的AI大语言模型驱动的同声传译、专业翻译助手,它拥有超精准的音频识别翻译能力,几乎零延迟的使用体验和支持多国语言可以让你带它走遍全球,无论你是留学生、商务人士、韩剧美剧爱好者,还是出国游玩、多国会议、跨国追星等等,都可以满足你所有需要同传的场景需求,线上线下通用,扫除语言障碍,让全世界的语言交流不再有国界。

讯飞智文

讯飞智文

一键生成PPT和Word,让学习生活更轻松

讯飞智文是一个利用 AI 技术的项目,能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定,还是学生群体的职业生涯规划、实习避坑指南,亦或是活动策划、旅游攻略等内容,它都能提供支持,帮助用户精准表达,轻松呈现各种信息。

AI办公办公工具AI工具讯飞智文AI在线生成PPTAI撰写助手多语种文档生成AI自动配图热门
下拉加载更多