IVF-Flat#

索引构建参数#

class cuvs.neighbors.ivf_flat.IndexParams( n_lists=1024, *, metric='sqeuclidean', metric_arg=2.0, kmeans_n_iters=20, kmeans_trainset_fraction=0.5, adaptive_centers=False, add_data_on_build=True, conservative_memory_allocation=False, )#

用于构建 IvfFlat 最近邻搜索索引的参数

参数:

n_listsint, default = 1024

粗量化器中使用的聚类数量。

metricstr, default = “sqeuclidean”

表示度量类型（metric type）的字符串。metric 的有效值包括：[“sqeuclidean”, “inner_product”, “euclidean”, “cosine”]，其中

sqeuclidean 是未进行平方根运算的欧几里得距离，即：distance(a,b) = sum_i (a_i - b_i)^2，

euclidean 是欧几里得距离

内积距离（inner product distance）定义为 distance(a, b) = sum_i a_i * b_i。

余弦距离（cosine distance）定义为 distance(a, b) = 1 - sum_i a_i * b_i / ( ||a||_2 * ||b||_2)。

kmeans_n_itersint, default = 20

在索引构建期间搜索 kmeans 中心点的迭代次数。默认设置通常是合适的，但对于较大的训练集比例（1000万+向量），可以减小此参数以缩短训练时间；对于较小的训练集比例（少量向量），可以增加此参数以提高召回率（recall）。

kmeans_trainset_fractionint, default = 0.5

如果 kmeans_trainset_fraction 小于 1，则数据集将被采样，仅使用 n_samples * kmeans_trainset_fraction 行数据进行训练。

add_data_on_buildbool, default = True

训练完粗粒度和细粒度量化器后，如果 add_data_on_build == True，我们将使用数据集填充索引；否则，索引将保持为空，可以使用 extend 方法向索引添加新向量。

adaptive_centersbool, default = False

默认情况下（adaptive_centers = False），聚类中心在 ivf_flat.build 中训练，并且在 ivf_flat.extend 中永远不会被修改。另一种行为（adaptive_centers = true）是在添加新数据时更新聚类中心。在这种情况下，index.centers() 总是对应聚类中数据的精确质心。这种行为的缺点是质心依赖于添加新数据的顺序（通过对添加数据的分类）；也就是说，index.centers() 会随着新添加数据的分布变化而“漂移”。

属性:

adaptive_centers
add_data_on_build
conservative_memory_allocation
kmeans_n_iters
kmeans_trainset_fraction
metric
metric_arg
n_lists

索引搜索参数#

class cuvs.neighbors.ivf_flat.SearchParams(n_probes=20, *)#

搜索 IVF-Flat 索引的补充参数

参数:

n_probes: int: 要搜索的聚类数量。

属性:

n_probes

索引#

class cuvs.neighbors.ivf_flat.Index#

IvfFlat 索引对象。此对象存储训练好的 IvfFlat 索引状态，可用于执行最近邻搜索。

属性:

trained

索引构建#

cuvs.neighbors.ivf_flat.build(IndexParams index_params, dataset, resources=None)[source]#

从数据集构建 IvfFlat 索引以进行高效搜索。

参数:

index_paramscuvs.neighbors.ivf_flat.IndexParams
dataset符合 CUDA 数组接口的矩阵，形状为 (n_samples, dim): 支持的数据类型（dtype）：[float, int8, uint8]
resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。: 如果未提供 Resource，将在函数内部分配 CUDA 资源并在函数退出前同步。如果提供了 Resource，则需要在访问输出之前通过调用 resources.sync() 显式同步。

返回值:

index: py:class:cuvs.neighbors.ivf_flat.Index

示例

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_flat
>>> 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 = ivf_flat.IndexParams(metric="sqeuclidean")
>>> index = ivf_flat.build(build_params, dataset)
>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
...                                        index, dataset,
...                                        k)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

索引搜索#

cuvs.neighbors.ivf_flat.search( SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None, filter=None, )[source]#

找到每个查询的 k 个最近邻。

参数:

search_paramspy:class:cuvs.neighbors.ivf_flat.SearchParams
indexpy:class:cuvs.neighbors.ivf_flat.Index: 训练好的 IvfFlat 索引。
queries符合 CUDA 数组接口的矩阵，形状为 (n_samples, dim): 支持的数据类型（dtype）：[float, int8, uint8]
kint: 最近邻的数量。
neighbors可选的符合 CUDA 数组接口的矩阵，形状为: (n_queries, k)，数据类型（dtype）为 int64_t。如果提供，最近邻索引将直接写入此处。（默认值为 None）
distances可选的符合 CUDA 数组接口的矩阵，形状为: (n_queries, k)。如果提供，到最近邻的距离将直接写入此处。（默认值为 None）
filter: 可选的 cuvs.neighbors.cuvsFilter 可用于过滤: 基于给定位集（bitset）的最近邻。（默认值为 None）
resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。: 如果未提供 Resource，将在函数内部分配 CUDA 资源并在函数退出前同步。如果提供了 Resource，则需要在访问输出之前通过调用 resources.sync() 显式同步。

示例

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_flat
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build the index
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset)
>>>
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> search_params = ivf_flat.SearchParams(n_probes=20)
>>>
>>> distances, neighbors = ivf_flat.search(search_params, index, queries,
...                                     k)

索引保存#

cuvs.neighbors.ivf_flat.save(filename, Index index, bool include_dataset=True, resources=None)[source]#

将索引保存到文件。

索引的保存/加载功能尚处于实验阶段。序列化格式可能会更改。

参数:

filenamestring: 文件名。
indexIndex: 训练好的 IVF-Flat 索引。
resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。: 如果未提供 Resource，将在函数内部分配 CUDA 资源并在函数退出前同步。如果提供了 Resource，则需要在访问输出之前通过调用 resources.sync() 显式同步。

示例

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_flat
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset)
>>> # Serialize and deserialize the ivf_flat index built
>>> ivf_flat.save("my_index.bin", index)
>>> index_loaded = ivf_flat.load("my_index.bin")

索引加载#

cuvs.neighbors.ivf_flat.load(filename, resources=None)[source]#

从文件加载索引。

索引的保存/加载功能尚处于实验阶段。序列化格式可能会更改，因此不保证能够加载使用先前版本 cuvs 保存的索引。

参数:

filenamestring: 文件名。
resources可选的 cuVS 资源句柄，用于复用 CUDA 资源。: 如果未提供 Resource，将在函数内部分配 CUDA 资源并在函数退出前同步。如果提供了 Resource，则需要在访问输出之前通过调用 resources.sync() 显式同步。

返回值:

indexIndex