cuVS: GPU上的向量搜索和聚类

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

RAFT为所有近似最近邻和聚类算法提供C++和Python API。

cuVS对不同语言提供越来越多的支持，包括C、C++、Python和Rust。我们将在未来为cuVS添加更多语言支持，但不会改进RAFT的语言支持。

一旦RAFT的所有近似最近邻和聚类算法都迁移到cuVS，RAFT的API将被弃用并最终完全移除。移除后，RAFT将成为一个轻量级的仅头文件库。在此期间，如果不需要额外的语言支持，使用RAFT也没有问题。

有用资源

代码示例：独立的代码示例。
API参考文档：API文档。
入门指南：RAFT入门。
构建和安装指南：安装和构建cuVS的说明。
RAPIDS社区：获取帮助、贡献和协作。
GitHub仓库：下载cuVS源代码。
问题追踪器：报告问题或请求功能。

什么是cuVS？

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

安装cuVS

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

Python	C/C++
`cuvs`	`libcuvs`, `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}
}