RAFT：用于向量搜索和其他的可重用加速函数和工具

[!IMPORTANT] RAFT中的向量搜索和聚类算法正在迁移到一个新的专门用于向量搜索的库，称为cuVS。在这次迁移期间，我们将继续支持RAFT中的向量搜索算法，但在RAPIDS 24.06（6月）发布后将不再更新它们。我们计划在RAPIDS 24.08（8月）发布时完成迁移。

内容

1. 有用资源
2. 什么是RAFT？
3. 用例
4. RAFT适合我吗？
5. 入门指南
6. 安装RAFT
7. 代码库结构和内容
8. 贡献
9. 参考文献

有用资源

• RAFT参考文档：API文档。
• RAFT入门指南：如何开始使用RAFT。
• 构建和安装 RAFT：安装和构建 RAFT 的说明。
• 示例笔记本：示例 jupyter 笔记本。
• RAPIDS 社区：获取帮助、贡献和协作。
• GitHub 仓库：下载 RAFT 源代码。
• 问题跟踪器：报告问题或请求功能。

什么是RAFT？

RAFT 包含机器学习和信息检索广泛使用的算法和基本原语。这些算法经过 CUDA 加速，形成构建高性能应用程序的基本块。

通过采用基于原语的算法开发方法，RAFT

• 缩短了算法构建时间
• 通过最大化项目间的重用减少维护负担，并且
• 集中核心可重用计算，允许未来的优化惠及使用它们的所有算法。

虽然不全面，但以下一般类别有助于总结 RAFT 中加速的功能

RAFT 是一个仅包含 C++ 头文件的模板库，可选的共享库可以

1. 加快常见模板类型的编译时间，并且
2. 提供主机可访问的 "运行时" API，无需 CUDA 编译器即可使用

除了是 C++ 库之外，RAFT 还提供了 2 个 Python 库

• pylibraft - 在 RAFT 的主机可访问的 "运行时" API 周围的轻量级 Python 包装器。
• raft-dask - 用于在 GPU 上使用 Dask 构建分布式算法的多节点多 GPU 通信基础设施。

用例

向量相似度搜索

RAFT 包含 GPU 上近似最近邻搜索 (ANNS) 算法的最先进实现，例如

• 暴力搜索。在不使用索引的情况下执行暴力最近邻搜索。
• IVF-Flat 和 IVF-PQ。使用倒排文件索引结构将内容映射到其位置。IVF-PQ 还使用产品量化来减少向量的内存使用。这些方法最初是由 FAISS 库普及的。
• CAGRA（Cuda Anns GRAph-based）。使用针对 GPU 优化的快速 ANNS 图构建和搜索实现。CAGRA 在大型批查询、单个查询和图构建时间方面优于最先进的 CPU 方法（即 HNSW）。

使用 RAFT ANNS 算法加速向量搜索的项目包括：Milvus、Redis 和 Faiss。

请参阅示例 Jupyter 笔记本 以开始使用 Python 中的 RAFT 进行向量搜索。

信息检索

RAFT 包含用于组合需要快速邻域计算的算法的可重用原语目录，例如

1. 计算向量之间的距离和计算核语法矩阵
2. 执行球半径查询以构建ε邻域
3. 将点聚类以分割空间，以便进行更小和更快的搜索
4. 从密集向量构建邻域“连通性”图

机器学习

RAFT的基本操作被多个RAPIDS库使用，包括cuML、cuGraph和cuOpt，以构建跨越大量不同应用的端到端机器学习算法，包括：

• 数据生成
• 模型评估
• 分类和回归
• 聚类
• 流形学习
• 降维。

RAFT还由流行的协同过滤库implicit用于推荐系统。

RAFT适合我吗？

RAFT包含用于加速应用程序和工作流程的低级基本操作。数据源提供者和应用程序开发者可能会发现特定的工具——例如ANN算法——非常有用。RAFT不打算由数据科学家直接用于发现和实验。有关数据科学工具，请参阅RAPIDS网站。

入门指南

RAPIDS内存管理器（RMM）

RAFT严重依赖RMM，它简化了全局配置不同分配策略的负担，这些策略适用于使用RMM的库。

多维数组

RAFT中的API接受mdspan多维数组视图来表示类似于Numpy Python库中的ndarray的高维数据。RAFT还包含相应的拥有mdarray结构，这简化了主机和设备（GPU）内存中多维数据的分配和管理。

mdarray在RMM上形成了一个便利层，并且可以在RAFT中使用许多不同的辅助函数构建

#include <raft/core/device_mdarray.hpp>

int n_rows = 10;
int n_cols = 10;

auto scalar = raft::make_device_scalar<float>(handle, 1.0);
auto vector = raft::make_device_vector<float>(handle, n_cols);
auto matrix = raft::make_device_matrix<float>(handle, n_rows, n_cols);

C++示例

RAFT中的大多数基本操作都接受一个raft::device_resources对象来管理创建成本高昂的资源，例如CUDA流、流池以及其他CUDA库（如cublas和cusolver）的句柄。

以下示例演示了创建一个RAFT句柄，并使用它与device_matrix和device_vector来分配内存、生成随机聚类和计算成对欧几里得距离

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdarray.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

auto input = raft::make_device_matrix<float, int>(handle, n_samples, n_features);
auto labels = raft::make_device_vector<int, int>(handle, n_samples);
auto output = raft::make_device_matrix<float, int>(handle, n_samples, n_samples);

raft::random::make_blobs(handle, input.view(), labels.view());

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);

还可以创建raft::device_mdspan视图，以使用原始指针和形状信息调用相同的API

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

float *input;
int *labels;
float *output;

...
// Allocate input, labels, and output pointers
...

auto input_view = raft::make_device_matrix_view(input, n_samples, n_features);
auto labels_view = raft::make_device_vector_view(labels, n_samples);
auto output_view = raft::make_device_matrix_view(output, n_samples, n_samples);

raft::random::make_blobs(handle, input_view, labels_view);

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);

Python示例

pylibraft包包含RAFT算法和基本操作的Python API。pylibraft通过非常轻量级、最小依赖性以及接受支持__cuda_array_interface__的任何对象（如CuPy的ndarray）而很好地集成到其他库中。此包中公开的RAFT算法数量在版本之间持续增长。

以下示例演示了计算CuPy数组之间的成对欧几里得距离。请注意，CuPy不是pylibraft的必需依赖项。

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)

output = pairwise_distance(in1, in2, metric="euclidean")

上述示例中的output数组为类型raft.common.device_ndarray，它支持cuda_array_interface，这使得它与其他也支持此接口的库（如CuPy、Numba、PyTorch和RAPIDS cuDF）兼容。CuPy支持DLPack，这还实现了从raft.common.device_ndarray到JAX和Tensorflow的零拷贝转换。

以下是将输出pylibraft.device_ndarray转换为CuPy数组的示例

cupy_array = cp.asarray(output)

并将其转换为PyTorch张量

import torch

torch_tensor = torch.as_tensor(output, device='cuda')

或将其转换为RAPIDS cuDF数据帧

cudf_dataframe = cudf.DataFrame(output)

当相应的库已在您的环境中安装并可用时，通过设置全局配置选项，所有RAFT计算API也可以自动执行此转换

import pylibraft.config
pylibraft.config.set_output_as("cupy")  # All compute APIs will return cupy arrays
pylibraft.config.set_output_as("torch") # All compute APIs will return torch tensors

您还可以指定一个接受pylibraft.common.device_ndarray的callable，并执行自定义转换。以下示例将所有输出转换为numpy数组

pylibraft.config.set_output_as(lambda device_ndarray: return device_ndarray.copy_to_host())

pylibraft还支持写入预分配的输出数组，因此任何支持__cuda_array_interface__的数组都可以就地写入

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
output = cp.empty((n_samples, n_samples), dtype=cp.float32)

pairwise_distance(in1, in2, out=output, metric="euclidean")

安装

RAFT的C++和Python库都可以通过Conda安装，Python库还可以通过Pip安装。

通过Conda安装C++和Python

安装RAFT最简单的方法是通过conda，并提供了几个包。

• libraft-headers C++头文件
• libraft（可选）包含预编译模板实例化和运行时API的C++共享库。
• pylibraft（可选）Python库
• raft-dask（可选）Python库，用于在Dask集群中部署使用RAFT raft::comms抽象层的多节点多GPU算法。
• raft-ann-bench（可选）用于轻松生成与其它最先进实现比较RAFT矢量搜索算法的基准测试工具。
• raft-ann-bench-cpu（可选）与上述类似的可重复基准测试工具，但不需要在机器上安装CUDA。可以用于具有竞争力CPU的环境中进行测试。

根据您的CUDA版本，使用以下命令通过conda安装所有RAFT包（将rapidsai替换为rapidsai-nightly以安装更最新但稳定性较低的夜间包）。首选使用mamba而非conda命令。

# for CUDA 11.8
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=11.8

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=12.5

注意，上述命令还将安装libraft-headers和libraft。

您还可以使用上述mamba命令单独安装conda包。例如，如果您想将RAFT的头文件和预编译的共享库安装到项目中

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia libraft libraft-headers cuda-version=12.5

如果您正在安装C++ API，请参阅使用libraft获取有关使用预编译共享库的更多信息。您还可以参考示例C++模板项目，其中包含您可以将其添加到项目中并针对安装的RAFT开发工件构建的CMake配置。

通过Pip安装Python

pylibraft和raft-dask都有实验性包，可以通过pip安装

pip install pylibraft-cu11 --extra-index-url=https://pypi.nvidia.com
pip install raft-dask-cu11 --extra-index-url=https://pypi.nvidia.com

这些包静态构建了RAFT的预编译实例化，因此C++头文件和预编译共享库将不会 readily 在您的代码中使用。

构建说明中包含有关从源构建RAFT并将其包含到下游项目中的更多详细信息。构建说明中还包括了从源构建RAFT C++和Python部分的更全面版本。

您可以在cpp/template目录中找到RAFT项目模板示例，它演示了如何使用RAFT构建新应用程序或将RAFT集成到现有的CMake项目中。

贡献

如果您有兴趣为RAFT项目做出贡献，请阅读我们的贡献指南。有关开发者指南、工作流程和原则的详细信息，请参阅开发者指南。

参考文献

在一般引用RAFT时，请考虑引用此GitHub项目。

@misc{rapidsai,
  title={Rapidsai/raft: RAFT contains fundamental widely-used algorithms and primitives for data science, Graph and machine learning.},
  url={https://github.com/rapidsai/raft},
  journal={GitHub},
  publisher={Nvidia RAPIDS},
  author={Rapidsai},
  year={2022}
}

如果引用稀疏成对距离API，请考虑使用以下bibtex

@article{nolet2021semiring,
  title={Semiring primitives for sparse neighborhood methods on the gpu},
  author={Nolet, Corey J and Gala, Divye and Raff, Edward and Eaton, Joe and Rees, Brad and Zedlewski, John and Oates, Tim},
  journal={arXiv preprint arXiv:2104.06357},
  year={2021}
}

如果引用单链接层次聚类API，请考虑以下bibtex

@misc{nolet2023cuslink,
      title={cuSLINK: Single-linkage Agglomerative Clustering on the GPU},
      author={Corey J. Nolet and Divye Gala and Alex Fender and Mahesh Doijade and Joe Eaton and Edward Raff and John Zedlewski and Brad Rees and Tim Oates},
      year={2023},
      eprint={2306.16354},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

如果引用CAGRA，请考虑以下bibtex

@misc{ootomo2023cagra,
      title={CAGRA: Highly Parallel Graph Construction and Approximate Nearest Neighbor Search for GPUs},
      author={Hiroyuki Ootomo and Akira Naruse and Corey Nolet and Ray Wang and Tamas Feher and Yong Wang},
      year={2024},
      series = {ICDE '24}
}

如果引用k选择例程，请考虑以下bibtex

@proceedings{10.1145/3581784,
    title = {Parallel Top-K Algorithms on GPU: A Comprehensive Study and New Methods},
    author={Jingrong Zhang, Akira Naruse, Xipeng Li, and Yong Wang},
    year = {2023},
    isbn = {9798400701092},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    location = {Denver, CO, USA},
    series = {SC '23}
}

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

@inproceedings{10.1145/3459637.3482344,
    author = {Wang, Hui and Zhao, Wan-Lei and Zeng, Xiangxiang and Yang, Jianye},
    title = {Fast K-NN Graph Construction by GPU Based NN-Descent},
    year = {2021},
    isbn = {9781450384469},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    url = {https://doi.org/10.1145/3459637.3482344},
    doi = {10.1145/3459637.3482344},
    abstract = {NN-Descent is a classic k-NN graph construction approach. It is still widely employed in machine learning, computer vision, and information retrieval tasks due to its efficiency and genericness. However, the current design only works well on CPU. In this paper, NN-Descent has been redesigned to adapt to the GPU architecture. A new graph update strategy called selective update is proposed. It reduces the data exchange between GPU cores and GPU global memory significantly, which is the processing bottleneck under GPU computation architecture. This redesign leads to full exploitation of the parallelism of the GPU hardware. In the meantime, the genericness, as well as the simplicity of NN-Descent, are well-preserved. Moreover, a procedure that allows to k-NN graph to be merged efficiently on GPU is proposed. It makes the construction of high-quality k-NN graphs for out-of-GPU-memory datasets tractable. Our approach is 100-250\texttimes{} faster than the single-thread NN-Descent and is 2.5-5\texttimes{} faster than the existing GPU-based approaches as we tested on million as well as billion scale datasets.},
    booktitle = {Proceedings of the 30th ACM International Conference on Information \& Knowledge Management},
    pages = {1929–1938},
    numpages = {10},
    keywords = {high-dimensional, nn-descent, gpu, k-nearest neighbor graph},
    location = {Virtual Event, Queensland, Australia},
    series = {CIKM '21}
}