IVF-PQ#

索引构建参数#

class cuvs.neighbors.ivf_pq.IndexParams(
n_lists=1024,
*,
metric='sqeuclidean',
metric_arg=2.0,
kmeans_n_iters=20,
kmeans_trainset_fraction=0.5,
pq_bits=8,
pq_dim=0,
codebook_kind='subspace',
force_random_rotation=False,
add_data_on_build=True,
conservative_memory_allocation=False,
max_train_points_per_pq_code=256,
)#

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

参数:
n_listsint, 默认值 = 1024

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

metricstr, 默认值 = “sqeuclidean”

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

  • sqeuclidean 是没有开平方根的欧几里得距离，即：距离(a,b) = sum_i (a_i - b_i)^2，

  • euclidean 是欧几里得距离

  • inner product distance 定义为距离(a, b) = sum_i a_i * b_i。

kmeans_n_itersint, 默认值 = 20

在索引构建期间搜索 kmeans 中心点的迭代次数。

kmeans_trainset_fractionint, 默认值 = 0.5

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

pq_bitsint, 默认值 = 8

量化后向量元素的位长度。

pq_dimint, 默认值 = 0

乘积量化后向量的维度。当为零时，使用启发式方法选择最优值。请注意，pq_dim * pq_bits 必须是 8 的倍数。提示：较小的“pq_dim”会导致索引大小减小，搜索性能提高，但召回率降低。如果“pq_bits”为 8，“pq_dim”可以设置为任意数字，但 8 的倍数对于良好性能是可取的。如果“pq_bits”不为 8，“pq_dim”应该是 8 的倍数。为了获得良好性能，最好“pq_dim”是 32 的倍数。理想情况下，“pq_dim”也应该是数据集维度的约数。

codebook_kindstring, 默认值 = “subspace”

有效值 [“subspace”, “cluster”]

force_random_rotationbool, 默认值 = False

即使 dim % pq_dim == 0，也对输入数据和查询应用随机旋转矩阵。注意：如果 dim 不是 pq_dim 的倍数，总是会对输入数据和查询应用随机旋转以将工作空间从 dim 转换到 rot_dim，rot_dim 可能略大于原始空间且是 pq_dim 的倍数 (rot_dim % pq_dim == 0)。但是，当 dim 是 pq_dim 的倍数时 (dim == rot_dim)，此转换不是必需的（因此无需添加“额外”的数据列/特征）。默认情况下，如果 dim == rot_dim，旋转转换将使用单位矩阵初始化。当 force_random_rotation == True 时，无论 dim 和 pq_dim 的值如何，都会生成一个随机正交转换矩阵。

add_data_on_buildbool, 默认值 = True

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

conservative_memory_allocationbool, 默认值 = True

默认情况下，算法会为单个聚类 (list_data) 分配比必需空间更多的空间。这允许分摊内存分配的成本并减少在重复调用 extend（扩展数据库）期间的数据复制次数。要禁用此行为并尽可能少地使用 GPU 内存来存储数据库，请将此标志设置为 True。

max_train_points_per_pq_codeint, 默认值 = 256

在 PQ 码本训练期间每个 PQ 码使用的最大数据点数量。每个 PQ 码使用更多数据点可能会提高 PQ 码本的质量，但也可能增加构建时间。该参数应用于两种 PQ 码本生成方法，即 PER_SUBSPACE 和 PER_CLUSTER。在这两种情况下，我们将使用 pq_book_size * max_train_points_per_pq_code 个训练点来训练每个码本。

属性:
add_data_on_build
codebook_kind
conservative_memory_allocation
force_random_rotation
kmeans_n_iters
kmeans_trainset_fraction
max_train_points_per_pq_code
metric
metric_arg
n_lists
pq_bits
pq_dim

索引搜索参数#

class cuvs.neighbors.ivf_pq.SearchParams(
n_probes=20,
*,
lut_dtype=np.float32,
internal_distance_dtype=np.float32,
)#

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

参数:
n_probes: int

要搜索的聚类数量。

lut_dtype: 默认值 = np.float32

在搜索时动态创建的查找表的（look up table）数据类型。使用低精度类型可以减少搜索时所需的共享内存量，因此即使对于维度较大的数据集，也可以使用快速共享内存内核。请注意，选择低精度类型时，召回率会略有下降。可能的值 [np.float32, np.float16, np.uint8]

internal_distance_dtype: 默认值 = np.float32

距离/相似度计算的存储数据类型。可能的值 [np.float32, np.float16]

属性:
internal_distance_dtype
lut_dtype
n_probes

索引#

class cuvs.neighbors.ivf_pq.Index#

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

属性:
trained

索引构建#

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

从数据集构建 IvfPq 索引，以实现高效搜索。

参数:
index_paramscuvs.neighbors.ivf_pq.IndexParams

如何构建索引的参数

dataset符合 CUDA 数组接口的矩阵，形状为 (n_samples, dim)

支持的 dtype [float, int8, uint8]

resources可选的 cuVS 资源句柄，用于重用 CUDA 资源。

如果未提供 Resources，CUDA 资源将在函数内部分配，并在函数退出前同步。如果提供了 resources，您需要通过调用 resources.sync() 来显式同步，然后才能访问输出。

返回:
index: cuvs.neighbors.ivf_pq.Index

示例

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

索引保存#

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

将索引保存到文件。

索引的保存/加载是实验性的。序列化格式可能会更改。

参数:
filenamestring

文件名。

indexIndex

经过训练的 IVF-PQ 索引。

resources可选的 cuVS 资源句柄，用于重用 CUDA 资源。

如果未提供 Resources，CUDA 资源将在函数内部分配，并在函数退出前同步。如果提供了 resources，您需要通过调用 resources.sync() 来显式同步，然后才能访问输出。

示例

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

索引加载#

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

从文件加载索引。

索引的保存/加载是实验性的。序列化格式可能会更改，因此不保证加载使用旧版本 cuvs 保存的索引能够正常工作。

参数:
filenamestring

文件名。

resources可选的 cuVS 资源句柄，用于重用 CUDA 资源。

如果未提供 Resources，CUDA 资源将在函数内部分配，并在函数退出前同步。如果提供了 resources，您需要通过调用 resources.sync() 来显式同步，然后才能访问输出。

返回:
indexIndex