cuVS Bench

cuVS bench 为各种 ANN 搜索实现提供了可复现的基准测试工具。它特别适合比较 GPU 实现以及比较 GPU 和 CPU。cuVS 的主要目标之一是针对各种重要的使用模式捕获理想的索引配置，以便结果可以在不同的硬件环境（例如本地和云端）上轻松重现。

此工具提供多项优势，包括：

1. 公平比较索引构建时间

2. 公平比较索引搜索吞吐量和/或延迟

3. 为不同召回率范围查找最佳参数设置

4. 轻松生成索引构建和搜索的一致样式图表

5. 剖析盲点和算法优化的潜力

6. 研究不同参数设置、索引构建时间和搜索性能之间的关系。

• 安装基准测试

  • Conda

  • Docker

• 运行基准测试

  • 端到端：小规模基准测试（<100 万至 1000 万）

  • 端到端：大规模基准测试（>1000 万向量）

  • 使用 Docker 容器运行

    • 在 GPU 上端到端运行

    • 手动在容器内运行脚本

  • 评估结果

• 创建和自定义数据集配置

  • 多 GPU 基准测试

• 添加新的索引算法

  • 实现与配置

  • 添加 Cmake 目标

安装基准测试

预编译的基准测试主要有两种分发方式：

• Conda 适用于不使用容器但想要易于安装和使用的 Python 包的用户。计划添加 Pip wheel 作为无法使用 conda 且不愿使用容器用户的替代方案。

• Docker 只需 docker 和 [NVIDIA docker](NVIDIA/nvidia-docker) 即可使用。提供了一个单独的 docker run 命令用于基本数据集基准测试，以及容器内 conda 解决方案的所有功能。

Conda

conda create --name cuvs_benchmarks
conda activate cuvs_benchmarks

# to install GPU package:
conda install -c rapidsai -c conda-forge -c nvidia cuvs-bench=<rapids_version> cuda-version=11.8*

# to install CPU package for usage in CPU-only systems:
conda install -c rapidsai -c conda-forge  cuvs-bench-cpu

通道 rapidsai 可以轻松替换为 rapidsai-nightly 如果需要夜间构建的基准测试。CPU 包目前允许运行 HNSW 基准测试。

请参阅构建说明从源构建基准测试。

Docker

我们为支持 GPU 的系统以及没有 GPU 的系统提供镜像。以下镜像可用：

• cuvs-bench：包含 GPU 和 CPU 基准测试，可运行所有支持的算法。根据需要下载百万级别数据集。最适合偏好较小容器大小的 GPU 系统用户。运行 GPU 算法需要 NVIDIA Container Toolkit，运行 CPU 算法则不需要。

• cuvs-bench-datasets：包含已预装百万级别数据集的 GPU 和 CPU 基准测试。最适合想要运行容器中已包含的多个百万级别数据集的用户。

• cuvs-bench-cpu：仅包含 CPU 基准测试，大小最小。最适合想要使用最小容器在没有 GPU 的系统上重现基准测试的用户。

夜间构建镜像位于 dockerhub。

以下命令拉取针对 Python 3.10、CUDA 12.5 和 cuVS 24.12 版本的夜间构建容器：

docker pull rapidsai/cuvs-bench:24.12a-cuda12.5-py3.10 # substitute cuvs-bench for the exact desired container.

CUDA 和 Python 版本可以更改为支持的值：- 支持的 CUDA 版本：11.8 和 12.5 - 支持的 Python 版本：3.10 和 3.11。

您也可以在 dockerhub 网站上查看确切的版本：- cuVS bench images - 预装百万级别数据集的 cuVS bench 镜像 - 仅 CPU 的 cuVS bench 镜像

注意：GPU 容器使用容器内的 CUDA Toolkit，唯一的要求是主机上安装了支持该版本的驱动程序。因此，例如，CUDA 11.8 容器可以在支持 CUDA 12.x 的驱动程序系统上运行。另请注意，在 docker 容器内使用 GPU 需要 Nvidia Container Toolkit 中的 Nvidia-Docker 运行时。

运行基准测试

端到端：小规模基准测试（<100 万至 1000 万）

以下步骤演示了如何下载、安装和运行 Yandex Deep-1B 数据集中 1000 万向量子集的基准测试。默认情况下，如果定义了 RAPIDS_DATASET_ROOT_DIR 环境变量，数据集将存储并使用该变量指示的文件夹；否则，将使用脚本调用位置下的 datasets 子文件夹。

# (1) prepare dataset.
python -m cuvs_bench.get_dataset --dataset deep-image-96-angular --normalize

# (2) build and search index
python -m cuvs_bench.run --dataset deep-image-96-inner --algorithms cuvs_cagra --batch-size 10 -k 10

# (3) export data
python -m cuvs_bench.run --data-export --dataset deep-image-96-inner

# (4) plot results
python -m cuvs_bench.plot --dataset deep-image-96-inner

数据集名称	训练行数	维度	测试行数	距离
deep-image-96-angular	1000 万	96	1 万	Angular
fashion-mnist-784-euclidean	6 万	784	1 万	欧氏距离
glove-50-angular	110 万	50	1 万	Angular
glove-100-angular	110 万	100	1 万	Angular
mnist-784-euclidean	6 万	784	1 万	欧氏距离
nytimes-256-angular	29 万	256	1 万	Angular
sift-128-euclidean	100 万	128	1 万	欧氏距离

上述所有数据集都包含具有 100 个邻居的地面真值测试数据集。因此，对于这些数据集，k 必须小于或等于 100。

端到端：大规模基准测试（>1000 万向量）

cuvs_bench.get_dataset 由于其大小，不能用于下载`十亿级别数据集`_。您应该改为使用我们的十亿级别数据集指南来下载和准备它们。下载十亿级别数据集后，以下提到的所有其他 python 命令均可正常工作。

要下载十亿级别数据集，请访问 big-ann-benchmarks

我们还提供了一个名为 wiki-all 的新数据集，包含 8800 万个 768 维向量。此数据集旨在对大规模现实世界的检索增强生成 (RAG)/LLM 嵌入大小进行基准测试。它还包含 100 万和 1000 万向量子集，用于小规模实验。有关更多信息和下载数据集，请参阅我们的Wiki-all 数据集指南。

以下步骤演示了如何下载、安装和运行 Yandex Deep-1B 数据集中 1 亿向量子集的基准测试。请注意，此类规模的数据集建议用于内存较大的 GPU，例如 A100 或 H100。

mkdir -p datasets/deep-1B
# (1) prepare dataset
# download manually "Ground Truth" file of "Yandex DEEP"
# suppose the file name is deep_new_groundtruth.public.10K.bin
python -m cuvs_bench.split_groundtruth --groundtruth datasets/deep-1B/deep_new_groundtruth.public.10K.bin
# two files 'groundtruth.neighbors.ibin' and 'groundtruth.distances.fbin' should be produced

# (2) build and search index
python -m cuvs_bench.run --dataset deep-1B --algorithms cuvs_cagra --batch-size 10 -k 10

# (3) export data
python -m cuvs_bench.run --data-export --dataset deep-1B

# (4) plot results
python -m cuvs_bench.plot --dataset deep-1B

python -m cuvs_bench.split_groundtruth 的用法如下：

使用 Docker 容器运行

使用 Docker 容器运行基准测试提供了两种方法。

在 GPU 上端到端运行

当未提供其他入口点时，端到端脚本将运行上述运行基准测试中的所有步骤。

对于支持 GPU 的系统，DATA_FOLDER 变量应该是您希望数据集存储在 $DATA_FOLDER/datasets 中、结果存储在 $DATA_FOLDER/result 中的本地文件夹（我们强烈建议 $DATA_FOLDER 是用于容器数据集和结果的专用文件夹）

export DATA_FOLDER=path/to/store/datasets/and/results
docker run --gpus all --rm -it -u $(id -u)                      \
    -v $DATA_FOLDER:/data/benchmarks                            \
    rapidsai/cuvs-bench:24.10a-cuda11.8-py3.10              \
    "--dataset deep-image-96-angular"                           \
    "--normalize"                                               \
    "--algorithms cuvs_cagra,cuvs_ivf_pq --batch-size 10 -k 10" \
    ""

上述命令的用法如下：

参数	描述
rapidsai/cuvs-bench:24.10a-cuda11.8-py3.10	要使用的镜像。可以是 cuvs-bench 或 cuvs-bench-datasets
"--dataset deep-image-96-angular"	数据集名称
"--normalize"	是否标准化数据集
"--algorithms cuvs_cagra,hnswlib --batch-size 10 -k 10"	传递给 run 脚本的参数，例如要进行基准测试的算法、批量大小和 k
""	将传递给 plot 脚本的附加（可选）参数。

*关于用户和文件权限的注意事项：* -u $(id -u) 标志允许容器内的用户与容器外用户的 uid 匹配，从而允许容器读写挂载的卷（由 $DATA_FOLDER 变量指示）。

在 CPU 上端到端运行

上述章节中的容器参数也可用于仅 CPU 的容器，可在未安装 GPU 的系统上使用。

*注意：* 镜像更改为 cuvs-bench-cpu 容器，且不再使用 --gpus all 参数

export DATA_FOLDER=path/to/store/datasets/and/results
docker run  --rm -it -u $(id -u)                  \
    -v $DATA_FOLDER:/data/benchmarks              \
    rapidsai/cuvs-bench-cpu:24.10a-py3.10     \
     "--dataset deep-image-96-angular"            \
     "--normalize"                                \
     "--algorithms hnswlib --batch-size 10 -k 10" \
     ""

手动在容器内运行脚本

所有 cuvs-bench 镜像都包含 Conda 包，因此可以直接通过登录容器本身来使用它们

export DATA_FOLDER=path/to/store/datasets/and/results
docker run --gpus all --rm -it -u $(id -u)          \
    --entrypoint /bin/bash                          \
    --workdir /data/benchmarks                      \
    -v $DATA_FOLDER:/data/benchmarks                \
    rapidsai/cuvs-bench:24.10a-cuda11.8-py3.10

这将使您进入容器内的命令行，cuvs-bench Python 包已准备就绪，如上述[运行基准测试](#running-the-benchmarks)部分所述

(base) root@00b068fbb862:/data/benchmarks# python -m cuvs_bench.get_dataset --dataset deep-image-96-angular --normalize

此外，容器可以以分离模式运行，没有任何问题。

评估结果

基准测试捕获了几种不同的测量结果。下表描述了索引构建基准测试的每种测量结果：

名称	描述
基准测试	唯一标识基准测试实例的名称
时间	训练索引所花费的实际时间
CPU	训练索引所花费的 CPU 时间
迭代次数	迭代次数（通常为 1）
GPU	构建所花费的 GPU 时间
index_size	用于训练索引的向量数量

下表描述了索引搜索基准测试的每种测量结果。最重要的测量结果是 Latency、items_per_second、end_to_end。

名称	描述
基准测试	唯一标识基准测试实例的名称
时间	单个迭代（批次）的实际时间除以线程数。
CPU	平均 CPU 时间（用户时间 + 系统时间）。这不包括空闲时间（等待 GPU 同步时也可能发生）。
迭代次数	总批次数。这等于 total_queries / n_queries。
GPU	单个批次的 GPU 延迟（秒）。在吞吐量模式下，这是多个线程的平均值。
延迟	单个批次的延迟（秒），根据实际时间计算。在吞吐量模式下，这是多个线程的平均值。
召回率	正确邻居与地面真值邻居的比例。请注意，仅当在数据集配置中指定了地面真值文件时，才会显示此列。
items_per_second	总吞吐量，即每秒查询数 (QPS)。这大约等于 total_queries / end_to_end。
k	每次迭代中查询的邻居数量
end_to_end	运行所有迭代的所有批次所花费的总时间
n_queries	每个批次中总的查询向量数量
total_queries	所有迭代中查询的向量总数（= iterations * n_queries）

请注意以下几点：- Time 和 end_to_end 使用 slightly different 方法测量。这就是为什么 end_to_end = Time * Iterations 仅是近似成立的。- 屏幕上显示的实际表格可能会略有不同，因为对于每种不同的组合进行基准测试时，也会显示超参数。- 召回率计算：每次测试处理的查询数量取决于迭代次数。因此，如果处理的邻居少于基准测试可用数量，召回率可能会显示 slight fluctuations。

创建和自定义数据集配置

单个配置通常会定义一组算法及其相关的索引和搜索参数，这些参数可以在不同数据集之间通用。我们使用 YAML 定义特定于数据集和特定于算法的配置。

CUVS 在 ${CUVS_HOME}/python/cuvs_bench/src/cuvs_bench/run/conf 中提供了一个默认的 datasets.yaml 文件，其中包含多个数据集的配置。以下是 sift-128-euclidean 数据集的简单示例条目：

- name: sift-128-euclidean
  base_file: sift-128-euclidean/base.fbin
  query_file: sift-128-euclidean/query.fbin
  groundtruth_neighbors_file: sift-128-euclidean/groundtruth.neighbors.ibin
  dims: 128
  distance: euclidean

cuvs-bench 支持的 ANN 算法的配置文件位于 ${CUVS_HOME}/python/cuvs_bench/cuvs_bench/config/algos 中。cuvs_cagra 算法配置如下所示：

name: cuvs_cagra
groups:
  base:
    build:
      graph_degree: [32, 64]
      intermediate_graph_degree: [64, 96]
      graph_build_algo: ["NN_DESCENT"]
    search:
      itopk: [32, 64, 128]

  large:
    build:
      graph_degree: [32, 64]
    search:
      itopk: [32, 64, 128]

可以通过为具有 base 组的算法创建自定义 YAML 文件来覆盖运行基准测试的默认参数。

上面的配置有两个字段：1. name - 定义正在指定参数的算法名称。2. groups - 定义具有特定参数集的运行组。每个组有助于为 build 和 search 创建所有超参数字段的笛卡尔积。

下表包含 cuVS 支持的所有算法。每个独特的算法都有自己的一组 build 和 search 设置。ANN 算法参数调优指南包含有关如何为每个支持的算法选择构建和搜索参数的详细说明。

库	算法
FAISS_GPU	faiss_gpu_flat, faiss_gpu_ivf_flat, faiss_gpu_ivf_pq
FAISS_CPU	faiss_cpu_flat, faiss_cpu_ivf_flat, faiss_cpu_ivf_pq
GGNN	ggnn
HNSWLIB	hnswlib
cuVS	cuvs_brute_force, cuvs_cagra, cuvs_ivf_flat, cuvs_ivf_pq, cuvs_cagra_hnswlib

多 GPU 基准测试

cuVS 实现了单节点多 GPU 版本的 IVF-Flat、IVF-PQ 和 CAGRA 索引。

索引类型	多 GPU 算法名称
IVF-Flat	cuvs_mg_ivf_flat
IVF-PQ	cuvs_mg_ivf_pq
CAGRA	cuvs_mg_cagra

添加新的索引算法

实现与配置

新算法的实现应该是一个 C++ 类，它继承 class ANN（定义在 cpp/bench/ann/src/ann.h 中）并实现所有纯虚函数。

此外，它应该定义两个用于构建和搜索参数的 struct`s 。搜索参数类应该继承 struct ANN<T>::AnnSearchParam。以 class HnswLib 为例，其定义如下：

基准测试程序原生使用 JSON 格式在配置文件中指定要构建的索引，以及构建和搜索参数。然而，JSON 配置文件过于冗长，不适合直接使用。相反，Python 脚本解析 YAML 并自动创建这些 json 文件。重要的是要认识到这些 json 对象与 build_param（其值为 JSON 对象）和 search_param（其值为 JSON 对象数组）的 yaml 对象对齐。以 HnswLib 的 json 配置为例，展示从 yaml 解析后的 json：

构建和搜索参数最终作为 json 对象传递给 C++ 层，用于每个参数配置的基准测试。以下代码展示了如何为 Hnswlib 解析这些参数：

1. 首先，添加两个函数用于将 JSON 对象解析为 struct BuildParam 和 struct SearchParam：

template<typename T>
void parse_build_param(const nlohmann::json& conf,
                       typename cuann::HnswLib<T>::BuildParam& param) {
  param.ef_construction = conf.at("efConstruction");
  param.M = conf.at("M");
  if (conf.contains("numThreads")) {
    param.num_threads = conf.at("numThreads");
  }
}

template<typename T>
void parse_search_param(const nlohmann::json& conf,
                        typename cuann::HnswLib<T>::SearchParam& param) {
  param.ef = conf.at("ef");
  if (conf.contains("numThreads")) {
    param.num_threads = conf.at("numThreads");
  }
}

2. 接下来，通过调用解析函数，在 create_algo()（在 cpp/bench/ann/ 中）和 `create_search_param() 函数中添加相应的 if 条件语句。 if 条件语句中的字符串文字必须与配置文件中 algo 的值相同。例如：

添加 Cmake 目标

在 cuvs/cpp/bench/ann/CMakeLists.txt 中，我们提供了一个 CMake 函数来配置新的基准测试目标，其签名如下：

要为 HNSWLIB 添加一个目标，我们将调用该函数如下：

ConfigureAnnBench(
  NAME HNSWLIB PATH bench/ann/src/hnswlib/hnswlib_benchmark.cpp INCLUDES
  ${CMAKE_CURRENT_BINARY_DIR}/_deps/hnswlib-src/hnswlib CXXFLAGS "${HNSW_CXX_FLAGS}"
)

这将创建一个名为 HNSWLIB_ANN_BENCH 的可执行文件，然后可以使用它来运行 HNSWLIB 基准测试。

在 algos.yaml 中添加一个新的条目，将算法名称映射到其二进制可执行文件，并指定算法是否需要 GPU 支持。

executable：指定构建/搜索索引的二进制文件名称。假定该文件位于 cuvs/cpp/build/ 中。requires_gpu：表示算法是否需要 GPU 才能运行。