CAGRA#
CAGRA 是一种基于图的最近邻算法，专为 GPU 加速而从头开始构建。CAGRA 在小批量和大批量搜索的索引构建和查询性能方面都表现出最先进的水平。
索引构建参数#
- class cuvs.neighbors.cagra.IndexParams(
- metric='sqeuclidean',
- *,
- intermediate_graph_degree=128,
- graph_degree=64,
- build_algo='ivf_pq',
- nn_descent_niter=20,
- compression=None,
构建 CAGRA 最近邻搜索索引的参数
- 参数:
- metric表示度量类型的字符串，默认为“sqeuclidean”
- metric 的有效值：[“sqeuclidean”, “inner_product”]，其中
sqeuclidean 是没有平方根运算的欧几里得距离，即：distance(a,b) = sum_i (a_i - b_i)^2
inner_product 距离定义为 distance(a, b) = sum_i a_i * b_i。
- intermediate_graph_degreeint 类型，默认为 128
- graph_degreeint 类型，默认为 64
- build_algo: 表示要使用的图构建算法的字符串，默认为“ivf_pq”
- algo 的有效值：[“ivf_pq”, “nn_descent”,
“iterative_cagra_search”]，其中
ivf_pq 将使用 IVF-PQ 算法构建 knn 图
nn_descent（实验性）将使用 NN-Descent 算法构建 knn 图。预计它通常会比 ivf_pq 更快。
iterative_cagra_search 将使用 CAGRA 的 search() 和 optimize() 迭代构建 knn 图
- compression: CompressionParams，可选
如果需要压缩，应为一个 CompressionParams 对象。如果为 None，则禁用压缩。
- 属性:
- build_algo
compressioncompression: object
- graph_degree
- intermediate_graph_degree
- metric
- nn_descent_niter
- compression#
compression: object
索引搜索参数#
- class cuvs.neighbors.cagra.SearchParams(
- max_queries=0,
- *,
- itopk_size=64,
- max_iterations=0,
- algo='auto',
- team_size=0,
- search_width=1,
- min_iterations=0,
- thread_block_size=0,
- hashmap_mode='auto',
- hashmap_min_bitlen=0,
- hashmap_max_fill_rate=0.5,
- num_random_samplings=1,
- rand_xor_mask=0x128394,
CAGRA 搜索参数
- 参数:
- max_queries: int 类型，默认为 0
同时搜索的最大查询数（批量大小）。为 0 时自动选择。
- itopk_size: int 类型，默认为 64
搜索过程中保留的中间搜索结果数量。这是调整精度和搜索速度之间权衡的主要旋钮。值越高，搜索精度越高。
- max_iterations: int 类型，默认为 0
搜索迭代次数上限。为 0 时自动选择。
- algo: 表示要使用的搜索算法的字符串，默认为“auto”
- algo 的有效值：[“auto”, “single_cta”, “multi_cta”]，其中
auto 将根据查询大小自动选择最佳值
当查询包含大量向量（例如 >10）时，single_cta 更好
当查询只包含少量向量时，multi_cta 更好
- team_size: int 类型，默认为 0
用于计算单个距离的线程数。4、8、16 或 32。
- search_width: int 类型，默认为 1
在每次迭代中选择作为搜索起点的图节点数。
- min_iterations: int 类型，默认为 0
搜索迭代次数下限。
- thread_block_size: int 类型，默认为 0
线程块大小。0、64、128、256、512、1024。为 0 时自动选择。
- hashmap_mode: 表示要使用的哈希映射类型的字符串。
通常最好让算法自动选择此值，默认为“auto”。hashmap_mode 的有效值：[“auto”, “small”, “hash”]，其中
auto 将根据 algo 自动选择最佳值
small 将使用带重置的小型共享内存哈希表。
hash 将使用全局内存中的单个哈希表。
- hashmap_min_bitlen: int 类型，默认为 0
哈希映射填充率上限。大于 0.1，小于 0.9。
- hashmap_max_fill_rate: float 类型，默认为 0.5
哈希映射填充率上限。大于 0.1，小于 0.9。
- num_random_samplings: int 类型，默认为 1
初始随机种子节点选择的迭代次数。1 或更多。
- rand_xor_mask: int 类型，默认为 0x128394
用于初始随机种子节点选择的位掩码。
- 属性:
- algo
- hashmap_max_fill_rate
- hashmap_min_bitlen
- hashmap_mode
- itopk_size
- max_iterations
- max_queries
- min_iterations
- num_random_samplings
- rand_xor_mask
- search_width
- team_size
- thread_block_size
索引#
- class cuvs.neighbors.cagra.Index#
CAGRA 索引对象。此对象存储经过训练的 CAGRA 索引状态，可用于执行最近邻搜索。
- 属性:
- dim
- trained
索引构建#
- cuvs.neighbors.cagra.build(IndexParams index_params, dataset, resources=None)[source]#
从数据集构建 CAGRA 索引以进行高效搜索。
构建过程包含两个不同的步骤——首先构建一个中间 knn 图，然后对其进行优化以创建最终图。index_params 对象控制这些图的节点度。
要求数据集和优化后的图都适合 GPU 内存。
- 支持以下距离度量
L2
内积
- 参数:
- index_paramsIndexParams 对象
- dataset符合 CUDA 数组接口的矩阵，形状为 (n_samples, dim)
支持的数据类型 [float, int8, uint8]
- resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。
如果没有提供资源，CUDA 资源将在此函数内部分配并在函数退出前同步。如果提供了资源，则需要在访问输出之前显式调用
resources.sync()进行同步。
- 返回:
- index: cuvs.cagra.Index
示例
>>> import cupy as cp >>> from cuvs.neighbors import cagra >>> n_samples = 50000 >>> n_features = 50 >>> n_queries = 1000 >>> k = 10 >>> dataset = cp.random.random_sample((n_samples, n_features), ... dtype=cp.float32) >>> build_params = cagra.IndexParams(metric="sqeuclidean") >>> index = cagra.build(build_params, dataset) >>> distances, neighbors = cagra.search(cagra.SearchParams(), ... index, dataset, ... k) >>> distances = cp.asarray(distances) >>> neighbors = cp.asarray(neighbors)
索引搜索#
- cuvs.neighbors.cagra.search(
- SearchParams search_params,
- Index index,
- queries,
- k,
- neighbors=None,
- distances=None,
- resources=None,
- filter=None,
找到每个查询的 k 个最近邻。
- 参数:
- search_paramsSearchParams
- indexIndex
已训练的 CAGRA 索引。
- queries符合 CUDA 数组接口的矩阵，形状为 (n_samples, dim)
支持的数据类型 [float, int8, uint8]
- kint 类型
邻居的数量。
- neighbors可选的符合 CUDA 数组接口的矩阵，形状为
(n_queries, k)，数据类型为 int64_t。如果提供，邻居索引将在此处就地写入。（默认为 None）
- distances可选的符合 CUDA 数组接口的矩阵，形状为
(n_queries, k)。如果提供，到邻居的距离将在此处就地写入。（默认为 None）
- filter: 可选的 cuvs.neighbors.cuvsFilter 可用于过滤
基于给定位集的邻居。
（默认为 None）
- resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。
如果没有提供资源，CUDA 资源将在此函数内部分配并在函数退出前同步。如果提供了资源，则需要在访问输出之前显式调用
resources.sync()进行同步。
示例
>>> import cupy as cp >>> from cuvs.neighbors import cagra >>> n_samples = 50000 >>> n_features = 50 >>> n_queries = 1000 >>> dataset = cp.random.random_sample((n_samples, n_features), ... dtype=cp.float32) >>> # Build index >>> index = cagra.build(cagra.IndexParams(), dataset) >>> # Search using the built index >>> queries = cp.random.random_sample((n_queries, n_features), ... dtype=cp.float32) >>> k = 10 >>> search_params = cagra.SearchParams( ... max_queries=100, ... itopk_size=64 ... ) >>> # Using a pooling allocator reduces overhead of temporary array >>> # creation during search. This is useful if multiple searches >>> # are performed with same query size. >>> distances, neighbors = cagra.search(search_params, index, queries, ... k) >>> neighbors = cp.asarray(neighbors) >>> distances = cp.asarray(distances)
索引保存#
- cuvs.neighbors.cagra.save(filename, Index index, bool include_dataset=True, resources=None)[source]#
将索引保存到文件。
索引的保存/加载是实验性的。序列化格式可能会发生变化。
- 参数:
- filename字符串类型
文件名。
- indexIndex
已训练的 CAGRA 索引。
- include_dataset布尔类型
是否将数据集与索引一起写入。在序列化索引中包含数据集会占用额外的磁盘空间，如果您磁盘上已经有数据集的副本，可能不需要这样做。如果此选项设置为 false，则在加载索引后必须调用
index.update_dataset(dataset)。- resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。
如果没有提供资源，CUDA 资源将在此函数内部分配并在函数退出前同步。如果提供了资源，则需要在访问输出之前显式调用
resources.sync()进行同步。
示例
>>> import cupy as cp >>> from cuvs.neighbors import cagra >>> n_samples = 50000 >>> n_features = 50 >>> dataset = cp.random.random_sample((n_samples, n_features), ... dtype=cp.float32) >>> # Build index >>> index = cagra.build(cagra.IndexParams(), dataset) >>> # Serialize and deserialize the cagra index built >>> cagra.save("my_index.bin", index) >>> index_loaded = cagra.load("my_index.bin")