IVF-PQ#

IVF-PQ 方法是一种 ANN 算法。与 IVF-Flat 类似，IVF-PQ 将点分成若干个簇（也由参数 n_lists 指定），然后搜索最近的簇来计算最近邻（也由参数 n_probes 指定），但它使用一种称为乘积量化（product quantization）的技术来压缩向量的大小。

#include <cuvs/neighbors/ivf_pq.hpp>

namespace cuvs::neighbors::ivf_pq

索引构建参数#

enum class codebook_gen

用于指定如何创建 PQ 码本的类型。

值

enumerator PER_SUBSPACE

enumerator PER_CLUSTER

struct index_params : public cuvs::neighbors::index_params#

#include <ivf_pq.hpp>

公有成员

uint32_t n_lists = 1024#

倒排列表（簇）的数量

提示：每个簇的向量数量（n_rows/n_lists）应约为 1,000 到 10,000。

uint32_t kmeans_n_iters = 20#: 搜索 kmeans 中心（索引构建）的迭代次数。

double kmeans_trainset_fraction = 0.5#: 在迭代 kmeans 构建期间使用的数据比例。

uint32_t pq_bits = 8#

PQ 压缩后向量元素的位长。

可能的值：[4, 5, 6, 7, 8]。

提示：“pq_bits”越小，索引大小越小，搜索性能越好，但召回率越低。

uint32_t pq_dim = 0#

PQ 压缩后向量的维度。当为零时，使用启发式方法选择最优值。

注意：pq_dim * pq_bits 必须是 8 的倍数。

提示：“pq_dim”越小，索引大小越小，搜索性能越好，但召回率越低。如果“pq_bits”为 8，则“pq_dim”可以设置为任何数字，但 8 的倍数对于获得良好性能是可取的。如果“pq_bits”不为 8，则“pq_dim”应为 8 的倍数。为了获得良好性能，“pq_dim”理想情况下应为 32 的倍数。最好，“pq_dim”也应该是数据集 dim 的除数。

codebook_gen codebook_kind = codebook_gen::PER_SUBSPACE #: PQ 码本的创建方式。

bool force_random_rotation = false#

即使 dim % pq_dim == 0，也对输入数据和查询应用随机正交旋转矩阵。

注意：如果 dim 不是 pq_dim 的倍数，则始终会对输入数据和查询应用随机旋转，将工作空间从 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 的值如何，都会生成一个随机正交变换矩阵。

bool conservative_memory_allocation = false#

默认情况下，算法为单个簇（list_data）分配的空间会比实际需要的多。这有助于分摊内存分配的成本，并减少在重复调用 extend（扩展数据库）时的数据复制次数。

另一种方式是保守分配行为；启用时，算法始终分配存储给定数量记录所需的最小内存量。如果您希望为数据库使用尽可能少的 GPU 内存，请将此标志设置为 true。

bool add_data_on_build = true#

是否将数据集内容添加到索引，即：

true 意味着索引会用数据集向量填充，并在调用 build 后即可搜索。
false 意味着 build 只训练底层模型（例如量化器或聚类），但索引会留空；之后需要对索引调用 extend 来填充它。

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

公有静态函数

static index_params from_dataset( raft::matrix_extent<int64_t> dataset, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded )#

根据输入数据集的形状创建 index_params。使用示例

using namespace cuvs::neighbors;
raft::resources res;
// create index_params for a [N. D] dataset and have InnerProduct as the distance metric
auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
ivf_pq::index_params index_params =
  ivf_pq::index_params::from_dataset(dataset.extents(), raft::distance::InnerProduct);
// modify/update index_params as needed
index_params.add_data_on_build = true;

索引搜索参数#

struct search_params : public cuvs::neighbors::search_params#

#include <ivf_pq.hpp>

公有成员

uint32_t n_probes = 20#: 要搜索的簇数量。

cudaDataType_t lut_dtype = CUDA_R_32F#

搜索时动态创建的查找表的数据类型。

可能的值：[CUDA_R_32F, CUDA_R_16F, CUDA_R_8U]

使用低精度类型可以减少搜索时所需的共享内存量，因此即使对于维度较大的数据集，也可以使用快速共享内存内核。请注意，选择低精度类型时，召回率会略有下降。

cudaDataType_t internal_distance_dtype = CUDA_R_32F#

搜索时计算的距离/相似度的存储数据类型。

可能的值：[CUDA_R_16F, CUDA_R_32F]

如果搜索时的性能瓶颈是设备内存访问，选择 FP16 将略微提高性能。

double preferred_shmem_carveout = 1.0#

SM 的统一内存/L1 缓存中用作共享内存的优选比例。

可能的值：[0.0 - 1.0]，表示 sharedMemPerMultiprocessor 的比例。

人们希望增加 carveout 以确保主搜索内核有良好的 GPU 占用率，但又不能太高，以便为 L1 缓存留出一些内存。请注意，此值仅被视为提示。此外，GPU 通常只允许固定的缓存配置集，因此提供的值会向上取整到最近的配置。请参考目标 GPU 架构的 NVIDIA 调优指南。

请注意，这是一个低级调优参数，如果调整不当，可能会对搜索性能产生严重负面影响。

索引#

template<typename IdxT> struct index : public cuvs::neighbors::index#

#include <ivf_pq.hpp>

IVF-PQ 索引。

在 IVF-PQ 索引中，数据库向量 y 通过两级量化进行近似

y = Q_1(y) + Q_2(y - Q_1(y))

第一级量化器 (Q_1) 将向量 y 映射到最近的簇中心。簇的数量是 n_lists。

第二级量化器对残差进行编码，并被定义为乘积量化器 [1]。

乘积量化器用一个 pq_dim 维向量对一个 dim 维向量进行编码。首先，我们将输入向量分割成 pq_dim 个子向量（表示为 u），其中每个 u 向量包含 y 的 pq_len 个不同分量

y_1, y_2, … y_{pq_len}, y_{pq_len+1}, … y_{2*pq_len}, … y_{dim-pq_len+1} … y_{dim} ___________________/ ____________________________/ ______________________/ u_1 u_2 u_{pq_dim}

然后每个子向量用单独的量化器 q_i 编码，结果被连接起来

Q_2(y) = q_1(u_1),q_2(u_2),…,q_{pq_dim}(u_pq_dim})

每个量化器 q_i 输出一个带有 pq_bit 位的代码。第二级量化器也由相应子空间中的 k-means 聚类定义：重现值是质心，重现值集合是码本。

当数据维度 dim 不是 pq_dim 的倍数时，特征空间会使用随机正交矩阵进行变换，使其具有 rot_dim = pq_dim * pq_len 维度（rot_dim >= dim）。

第二级量化器针对每个子空间或每个簇进行训练：(a) codebook_gen::PER_SUBSPACE：创建 pq_dim 个第二级量化器 - 每个沿特征的数据切片一个；(b) codebook_gen::PER_CLUSTER：创建 n_lists 个第二级量化器 - 每个第一级簇一个。无论哪种情况，质心都是通过将数据解释为具有 pq_len 维度的 k-means 聚类再次找到的。

[1] 用于最近邻搜索的乘积量化 Herve Jegou, Matthijs Douze, Cordelia Schmid

模板参数:: IdxT – 源数据集中索引的类型

公有函数

index(raft::resources const &handle)#

构造一个空索引。

构造一个空索引。该索引需要通过 build 进行训练，或者通过 deserialize 从保存的副本加载。

index( raft::resources const &handle, cuvs::distance::DistanceType metric, codebook_gen codebook_kind, uint32_t n_lists, uint32_t dim, uint32_t pq_bits = 8, uint32_t pq_dim = 0, bool conservative_memory_allocation = false )#: 构造一个空索引。需要先训练，然后填充。

index( raft::resources const &handle, const index_params &params, uint32_t dim )#: 构造一个空索引。需要先训练，然后填充。

IdxT size() const noexcept#: 索引的总长度。

uint32_t dim() const noexcept#: 输入数据的维度。

uint32_t dim_ext() const noexcept#: 簇中心的维度：输入数据 dim 扩展了向量范数并填充到 8 个元素。

uint32_t rot_dim() const noexcept#: 经过 PQ 处理转换后的数据维度（旋转并增强，使其成为 pq_dim 的倍数）。

uint32_t pq_bits() const noexcept#: PQ 压缩后编码向量元素的位长。

uint32_t pq_dim() const noexcept#: PQ 压缩后编码向量的维度。

uint32_t pq_len() const noexcept#: 子空间的维度，即映射到子空间的向量分量数量

uint32_t pq_book_size() const noexcept#: PQ 码本中的向量数量（1 << pq_bits）。

cuvs::distance::DistanceType metric() const noexcept#: 用于聚类的距离度量。

codebook_gen codebook_kind() const noexcept#: PQ 码本的创建方式。

uint32_t n_lists() const noexcept#: 簇/倒排列表数量（第一级量化）。

bool conservative_memory_allocation() const noexcept#: 扩展列表（簇）数据时是否使用保守内存分配（参见 index_params.conservative_memory_allocation）。

raft::device_mdspan<float, pq_centers_extents, raft::row_major> pq_centers( ) noexcept#

PQ 簇中心

codebook_gen::PER_SUBSPACE: [pq_dim , pq_len, pq_book_size]
codebook_gen::PER_CLUSTER: [n_lists, pq_len, pq_book_size]

std::vector<std::shared_ptr<list_data<IdxT>>> &lists( ) noexcept#: 列表数据和索引。

raft::device_vector_view<uint8_t*, uint32_t, raft::row_major> data_ptrs( ) noexcept#: 倒排列表（簇）数据的指针 [n_lists]。

raft::device_vector_view<IdxT*, uint32_t, raft::row_major> inds_ptrs( ) noexcept#: 倒排列表（簇）索引的指针 [n_lists]。

raft::device_matrix_view<float, uint32_t, raft::row_major> rotation_matrix( ) noexcept#: 变换矩阵（原始空间 -> 旋转填充空间）[rot_dim, dim]

raft::host_vector_view<IdxT, uint32_t, raft::row_major> accum_sorted_sizes( ) noexcept#

累积的列表大小，按降序排序 [n_lists + 1]。最后一个值包含索引的总长度。索引零处的值始终为零。

也就是说，此 span 的内容就像 list_sizes 经过排序然后累积一样。

此 span 在搜索期间用于估计工作空间的最大大小。

raft::device_vector_view<uint32_t, uint32_t, raft::row_major> list_sizes( ) noexcept#: 列表的大小 [n_lists]。

raft::device_matrix_view<float, uint32_t, raft::row_major> centers( ) noexcept#: 原始空间中与列表对应的簇中心 [n_lists, dim_ext]

raft::device_matrix_view<float, uint32_t, raft::row_major> centers_rot( ) noexcept#: 旋转空间中与列表对应的簇中心 [n_lists, rot_dim]

uint32_t get_list_size_in_bytes(uint32_t label)#

使用列表范围获取特定 IVF 列表的字节大小。使用示例

raft::resources res;
// use default index params
ivf_pq::index_params index_params;
// extend the IVF lists while building the index
index_params.add_data_on_build = true;
// create and fill the index from a [N, D] dataset
auto index = cuvs::neighbors::ivf_pq::build<int64_t>(res, index_params, dataset, N, D);
// Fetch the size of the fourth list
uint32_t size = index.get_list_size_in_bytes(3);

参数:: label – [in] 列表 ID

索引构建#

cuvs::neighbors::ivf_pq::index<int64_t> build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::device_matrix_view<const float, int64_t, raft::row_major> dataset )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的设备矩阵视图 [n_rows, dim]

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::device_matrix_view<const float, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的 raft::device_matrix_view [n_rows, dim]
idx – [out] 对 ivf_pq::index 的引用

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的设备矩阵视图 [n_rows, dim]

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::device_matrix_view<const half, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的 raft::device_matrix_view [n_rows, dim]
idx – [out] 对 ivf_pq::index 的引用

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的设备矩阵视图 [n_rows, dim]

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的 raft::device_matrix_view [n_rows, dim]
idx – [out] 对 ivf_pq::index 的引用

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的设备矩阵视图 [n_rows, dim]

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [in] 指向行主序矩阵的 raft::device_matrix_view [n_rows, dim]
idx – [out] 对 ivf_pq::index 的引用

cuvs::neighbors::ivf_pq::index<int64_t> build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::host_matrix_view<const float, int64_t, raft::row_major> dataset )#

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 host_matrix_view

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::host_matrix_view<const float, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 raft::host_matrix_view
idx – [out] 对 ivf_pq::index 的引用

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 host_matrix_view

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::host_matrix_view<const half, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 raft::host_matrix_view
idx – [out] 对 ivf_pq::index 的引用

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

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 host_matrix_view

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 raft::host_matrix_view
idx – [out] 对 ivf_pq::index 的引用

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(handle, index_params, dataset);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 host_matrix_view

返回:

构建的 ivf-pq 索引

void build( raft::resources const &handle, const cuvs::neighbors::ivf_pq::index_params &index_params, raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

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

注意，如果 index_params.add_data_on_build 设置为 true，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping. This is only applicable if index_params.add_data_on_build is set to true
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// create and fill the index from a [N, D] dataset
ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
ivf_pq::build(handle, index_params, dataset, index);

参数:

handle – [in]
index_params – 配置索引构建
dataset – [输入] 指向行主序矩阵 [n_rows, dim] 的 raft::host_matrix_view
idx – [out] 对 ivf_pq::index 的引用

索引扩展#

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::device_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::device_matrix_view<const half, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::device_matrix_view<const half, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// fill the index with the data
std::optional<raft::device_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 device matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 device vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::host_matrix_view<const half, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::host_matrix_view<const half, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

cuvs::neighbors::ivf_pq::index<int64_t> extend( raft::resources const &handle, raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, const cuvs::neighbors::ivf_pq::index<int64_t> &idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

void extend( raft::resources const &handle, raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors, std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices, cuvs::neighbors::ivf_pq::index<int64_t> *idx )#

使用新数据扩展索引。

注意，用户可以在输入 raft::resource 中设置一个至少包含一个流的流池，以启用 kernel 和复制的重叠。

使用示例

using namespace cuvs::neighbors;
ivf_pq::index_params index_params;
index_params.add_data_on_build = false;      // don't populate index on build
index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
// train the index from a [N, D] dataset
auto index_empty = ivf_pq::build(handle, index_params, dataset);
// optional: create a stream pool with at least one stream to enable kernel and copy
// overlapping
raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(1));
// fill the index with the data
std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
ivf_pq::extend(handle, new_vectors, no_op, &index_empty);

参数:

handle – [in]
new_vectors – [输入] 指向行主序矩阵 [n_rows, idx.dim()] 的 host matrix view
new_indices – [输入] 指向索引向量 [n_rows] 的 host vector view。如果原始索引为空 (idx.size() == 0)，您可以在此处传递 std::nullopt 以表示连续范围 [0...n_rows)。
idx – [输入/输出]

索引搜索#

void search( raft::resources const &handle, const cuvs::neighbors::ivf_pq::search_params &search_params, const cuvs::neighbors::ivf_pq::index<int64_t> &index, raft::device_matrix_view<const float, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances, const cuvs::neighbors::filtering::base_filter &sample_filter = cuvs::neighbors::filtering::none_sample_filter{} )#

使用构建好的索引搜索 ANN。

有关用法示例，请参阅 ivf_pq::build 文档。

注意，此函数需要一个临时缓冲区用于存储 cuda kernel 调用之间的中间结果，这可能导致不期望的内存分配和速度下降。为了缓解此问题，您可以传递一个内存池资源或足够大的预分配内存资源，以减少或完全消除在 search 内部发生的分配。临时缓冲区的确切大小取决于多个因素，并且是实现细节。然而，您可以安全地为内存池指定一个较小的初始大小，这样在首次调用 search 期间只会发生少量分配来增长它。

...
// use default search parameters
ivf_pq::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_pq::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_pq::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_pq::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

参数:

handle – [in]
search_params – 配置搜索
index – [输入] 已构建的 ivf-pq 索引
queries – [输入] 指向行主序矩阵 [n_queries, index->dim()] 的 device matrix view
neighbors – [输出] 指向源数据集中近邻索引的 device matrix view [n_queries, k]
distances – [输出] 指向选定近邻距离的 device matrix view [n_queries, k]
sample_filter – [输入] 一个可选的设备过滤器函数对象，用于批准给定查询的样本。（none_sample_filter 表示不进行过滤）

void search( raft::resources const &handle, const cuvs::neighbors::ivf_pq::search_params &search_params, const cuvs::neighbors::ivf_pq::index<int64_t> &index, raft::device_matrix_view<const half, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances, const cuvs::neighbors::filtering::base_filter &sample_filter = cuvs::neighbors::filtering::none_sample_filter{} )#

使用构建好的索引搜索 ANN。

有关用法示例，请参阅 ivf_pq::build 文档。

...
// use default search parameters
ivf_pq::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_pq::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_pq::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_pq::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

参数:

handle – [in]
search_params – 配置搜索
index – [输入] 已构建的 ivf-pq 索引
queries – [输入] 指向行主序矩阵 [n_queries, index->dim()] 的 device matrix view
neighbors – [输出] 指向源数据集中近邻索引的 device matrix view [n_queries, k]
distances – [输出] 指向选定近邻距离的 device matrix view [n_queries, k]
sample_filter – [输入] 一个可选的设备过滤器函数对象，用于批准给定查询的样本。（none_sample_filter 表示不进行过滤）

void search( raft::resources const &handle, const cuvs::neighbors::ivf_pq::search_params &search_params, const cuvs::neighbors::ivf_pq::index<int64_t> &index, raft::device_matrix_view<const int8_t, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances, const cuvs::neighbors::filtering::base_filter &sample_filter = cuvs::neighbors::filtering::none_sample_filter{} )#

使用构建好的索引搜索 ANN。

有关用法示例，请参阅 ivf_pq::build 文档。

...
// use default search parameters
ivf_pq::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_pq::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_pq::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_pq::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

参数:

handle – [in]
search_params – 配置搜索
index – [输入] 已构建的 ivf-pq 索引
queries – [输入] 指向行主序矩阵 [n_queries, index->dim()] 的 device matrix view
neighbors – [输出] 指向源数据集中近邻索引的 device matrix view [n_queries, k]
distances – [输出] 指向选定近邻距离的 device matrix view [n_queries, k]
sample_filter – [输入] 一个可选的设备过滤器函数对象，用于批准给定查询的样本。（none_sample_filter 表示不进行过滤）

void search( raft::resources const &handle, const cuvs::neighbors::ivf_pq::search_params &search_params, const cuvs::neighbors::ivf_pq::index<int64_t> &index, raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> queries, raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors, raft::device_matrix_view<float, int64_t, raft::row_major> distances, const cuvs::neighbors::filtering::base_filter &sample_filter = cuvs::neighbors::filtering::none_sample_filter{} )#

使用构建好的索引搜索 ANN。

有关用法示例，请参阅 ivf_pq::build 文档。

...
// use default search parameters
ivf_pq::search_params search_params;
// Use the same allocator across multiple searches to reduce the number of
// cuda memory allocations
ivf_pq::search(handle, search_params, index, queries1, out_inds1, out_dists1);
ivf_pq::search(handle, search_params, index, queries2, out_inds2, out_dists2);
ivf_pq::search(handle, search_params, index, queries3, out_inds3, out_dists3);
...

参数:

handle – [in]
search_params – 配置搜索
index – [输入] 已构建的 ivf-pq 索引
queries – [输入] 指向行主序矩阵 [n_queries, index->dim()] 的 device matrix view
neighbors – [输出] 指向源数据集中近邻索引的 device matrix view [n_queries, k]
distances – [输出] 指向选定近邻距离的 device matrix view [n_queries, k]
sample_filter – [输入] 一个可选的设备过滤器函数对象，用于批准给定查询的样本。（none_sample_filter 表示不进行过滤）

索引序列化#

void serialize( raft::resources const &handle, std::ostream &os, const cuvs::neighbors::ivf_pq::index<int64_t> &index )#

将索引写入输出流

#include <raft/core/resources.hpp>

raft::resources handle;

// create an output stream
std::ostream os(std::cout.rdbuf());
// create an index with `auto index = ivf_pq::build(...);`
cuvs::neighbors::ivf_pq::serialize(handle, os, index);

参数:

handle – [输入] raft 句柄
os – [输入] 输出流
index – [输入] IVF-PQ 索引

void serialize( raft::resources const &handle, const std::string &filename, const cuvs::neighbors::ivf_pq::index<int64_t> &index )#

将索引保存到文件。

#include <raft/core/resources.hpp>

raft::resources handle;

// create a string with a filepath
std::string filename("/path/to/index");
// create an index with `auto index = ivf_pq::build(...);`
cuvs::neighbors::ivf_pq::serialize(handle, filename, index);

参数:

handle – [输入] raft 句柄
filename – [输入] 用于保存索引的文件名
index – [输入] IVF-PQ 索引

void deserialize( raft::resources const &handle, std::istream &str, cuvs::neighbors::ivf_pq::index<int64_t> *index )#

从输入流加载索引

#include <raft/core/resources.hpp>

raft::resources handle;

// create an input stream
std::istream is(std::cin.rdbuf());

using IdxT = int64_t; // type of the index
// create an empty index
cuvs::neighbors::ivf_pq::index<IdxT> index(handle);

cuvs::neighbors::ivf_pq::deserialize(handle, is, index);

参数:

handle – [输入] raft 句柄
str – [输入] 存储索引的文件名
index – [输出] IVF-PQ 索引

void deserialize( raft::resources const &handle, const std::string &filename, cuvs::neighbors::ivf_pq::index<int64_t> *index )#

从文件加载索引。

#include <raft/core/resources.hpp>

raft::resources handle;

// create a string with a filepath
std::string filename("/path/to/index");
using IdxT = int64_t; // type of the index
// create an empty index
ivf_pq::index<IdxT> index(handle);

cuvs::neighbors::ivf_pq::deserialize(handle, filename, &index);

参数:

handle – [输入] raft 句柄
filename – [输入] 存储索引的文件名
index – [输出] IVF-PQ 索引

辅助方法#

用于操作 IVF-PQ 索引底层数据、解包记录以及将 PQ 编码写入现有 IVF 列表的额外辅助函数。

namespace cuvs::neighbors::ivf_pq::helpers

void unpack( raft::resources const &res, raft::device_mdspan<const uint8_t, list_spec<uint32_t, uint32_t>::list_extents, raft::row_major> list_data, uint32_t pq_bits, uint32_t offset, raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> codes )#

解压压缩索引中从给定 offset 开始的单个列表（簇）的 n_take 个连续记录。

位压缩被移除，这意味着输出将具有 pq_dim 维向量（每字节一个编码，而不是 ceildiv(pq_dim * pq_bits, 8) 字节的 pq 编码）。

使用示例

auto list_data = index.lists()[label]->data.view();
// allocate the buffer for the output
uint32_t n_take = 4;
auto codes = raft::make_device_matrix<uint8_t>(res, n_take, index.pq_dim());
uint32_t offset = 0;
// unpack n_take elements from the list
ivf_pq::helpers::codepacker::unpack(res, list_data, index.pq_bits(), offset, codes.view());

参数:

res – [输入] raft 资源
list_data – 要读取的数据块
pq_bits – [输入] 编码向量元素的位长度
offset – [输入] 在列表中跳过多少记录。
codes – [输出] 目标缓冲区 [n_take, index.pq_dim()]。长度 n_take 定义要解包的记录数，它必须小于列表大小。

void unpack_contiguous( raft::resources const &res, raft::device_mdspan<const uint8_t, list_spec<uint32_t, uint32_t>::list_extents, raft::row_major> list_data, uint32_t pq_bits, uint32_t offset, uint32_t n_rows, uint32_t pq_dim, uint8_t *codes )#

解压压缩索引中从给定 offset 开始的单个列表（簇）的 n_rows 个连续记录。单个向量的输出编码是连续的，未扩展到每字节一个编码，这意味着每个 PQ 编码向量的输出占 ceildiv(pq_dim * pq_bits, 8) 字节。

使用示例

raft::resources res;
auto list_data = index.lists()[label]->data.view();
// allocate the buffer for the output
uint32_t n_rows = 4;
auto codes = raft::make_device_matrix<uint8_t>(
  res, n_rows, raft::ceildiv(index.pq_dim() * index.pq_bits(), 8));
uint32_t offset = 0;
// unpack n_rows elements from the list
ivf_pq::helpers::codepacker::unpack_contiguous(
  res, list_data, index.pq_bits(), offset, n_rows, index.pq_dim(), codes.data_handle());

参数:

res – [输入] raft 资源
list_data – 要读取的数据块
pq_bits – [输入] 编码向量元素的位长度
offset – [输入] 在列表中跳过多少记录。
n_rows – [输入] 要解包的记录数
pq_dim – [输入] PQ 压缩记录的维度
codes – [输出] 目标缓冲区 [n_rows, ceildiv(pq_dim * pq_bits, 8)]。长度 n_rows 定义要解包的记录数，它必须小于列表大小。

void pack( raft::resources const &res, raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> codes, uint32_t pq_bits, uint32_t offset, raft::device_mdspan<uint8_t, list_spec<uint32_t, uint32_t>::list_extents, raft::row_major> list_data )#

将平坦的 PQ 编码写入现有列表，使用给定的偏移量。

注意：此处不发生内存分配；列表必须能容纳数据 (offset + n_vec)。

使用示例

auto list_data  = index.lists()[label]->data.view();
// allocate the buffer for the input codes
auto codes = raft::make_device_matrix<uint8_t>(res, n_vec, index.pq_dim());
... prepare n_vecs to pack into the list in codes ...
// write codes into the list starting from the 42nd position
ivf_pq::helpers::codepacker::pack(
    res, make_const_mdspan(codes.view()), index.pq_bits(), 42, list_data);

参数:

res – [输入] raft 资源
codes – [输入] 平坦的 PQ 编码，每字节一个编码 [n_vec, pq_dim]
pq_bits – [输入] 编码向量元素的位长度
offset – [输入] 在将数据写入列表之前跳过多少记录
list_data – 要写入的数据块

void pack_contiguous( raft::resources const &res, const uint8_t *codes, uint32_t n_rows, uint32_t pq_dim, uint32_t pq_bits, uint32_t offset, raft::device_mdspan<uint8_t, list_spec<uint32_t, uint32_t>::list_extents, raft::row_major> list_data )#

将平坦的 PQ 编码写入现有列表，使用给定的偏移量。单个向量的输入编码是连续的（未扩展到每字节一个编码）。

注意：此处不发生内存分配；要打包向量的列表必须能容纳偏移量 + n_rows 条记录。

使用示例

raft::resources res;
auto list_data  = index.lists()[label]->data.view();
// allocate the buffer for the input codes
auto codes = raft::make_device_matrix<uint8_t>(
  res, n_rows, raft::ceildiv(index.pq_dim() * index.pq_bits(), 8));
... prepare compressed vectors to pack into the list in codes ...
// write codes into the list starting from the 42nd position. If the current size of the list
// is greater than 42, this will overwrite the codes starting at this offset.
ivf_pq::helpers::codepacker::pack_contiguous(
  res, codes.data_handle(), n_rows, index.pq_dim(), index.pq_bits(), 42, list_data);

参数:

res – [输入] raft 资源
codes – [输入] 平坦连续的 PQ 编码 [n_vec, ceildiv(pq_dim * pq_bits, 8)]
n_rows – [输入] 记录数
pq_dim – [输入]
pq_bits – [输入] 编码向量元素的位长度
offset – [输入] 在将数据写入列表之前跳过多少记录
list_data – 要写入的数据块

void pack_list_data( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> codes, uint32_t label, uint32_t offset )#

将平坦的 PQ 编码写入现有列表，使用给定的偏移量。

列表由其标签标识。

注意：此处不发生内存分配；列表必须能容纳数据 (offset + n_vec)。

使用示例

// We will write into the 137th cluster
uint32_t label = 137;
// allocate the buffer for the input codes
auto codes = raft::make_device_matrix<const uint8_t>(res, n_vec, index.pq_dim());
... prepare n_vecs to pack into the list in codes ...
// write codes into the list starting from the 42nd position
ivf_pq::helpers::codepacker::pack_list_data(res, &index, codes_to_pack, label, 42);

参数:

res – [输入] raft 资源
index – [输入/输出] IVF-PQ 索引。
codes – [输入] 平坦的 PQ 编码，每字节一个编码 [n_rows, pq_dim]
label – [输入] 要写入的列表（簇）的 ID。
offset – [输入] 在将数据写入列表之前跳过多少记录

void pack_contiguous_list_data( raft::resources const &res, index<int64_t> *index, uint8_t *codes, uint32_t n_rows, uint32_t label, uint32_t offset )#

将平坦的 PQ 编码写入现有列表，使用给定的偏移量。当输入向量经过 PQ 编码且未扩展到每字节一个编码时，请使用此函数。

列表由其标签标识。

注意：此处不发生内存分配；要打包向量的列表必须能容纳偏移量

n_rows 行。

使用示例

using namespace cuvs::neighbors;
raft::resources res;
// use default index parameters
ivf_pq::index_params index_params;
// create and fill the index from a [N, D] dataset
auto index = ivf_pq::build(res, index_params, dataset, N, D);
// allocate the buffer for n_rows input codes. Each vector occupies
// raft::ceildiv(index.pq_dim() * index.pq_bits(), 8) bytes because
// codes are compressed and without gaps.
auto codes = raft::make_device_matrix<const uint8_t>(
  res, n_rows, raft::ceildiv(index.pq_dim() * index.pq_bits(), 8));
... prepare the compressed vectors to pack into the list in codes ...
// the first n_rows codes in the fourth IVF list are to be overwritten.
uint32_t label = 3;
// write codes into the list starting from the 0th position
ivf_pq::helpers::codepacker::pack_contiguous_list_data(
  res, &index, codes.data_handle(), n_rows, label, 0);

参数:

res – [输入] raft 资源
index – [输入/输出] 指向 IVF-PQ 索引的指针
codes – [输入] 平坦连续的 PQ 编码 [n_rows, ceildiv(pq_dim * pq_bits, 8)]
n_rows – [输入] 要打包的记录数
label – [输入] 要写入的列表（簇）的 ID。
offset – [输入] 在将数据写入列表之前跳过多少记录

void unpack_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_codes, uint32_t label, uint32_t offset )#

解压压缩索引中从给定 offset 开始的单个列表（簇）的 n_take 个连续记录，每字节一个编码（与 pq_bits 无关）。

使用示例

  // We will unpack the fourth cluster
  uint32_t label = 3;
  // Get the list size
  uint32_t list_size = 0;
  raft::copy(&list_size, index.list_sizes().data_handle() + label, 1,
resource::get_cuda_stream(res)); resource::sync_stream(res);
  // allocate the buffer for the output
  auto codes = raft::make_device_matrix<float>(res, list_size, index.pq_dim());
  // unpack the whole list
  ivf_pq::helpers::codepacker::unpack_list_data(res, index, codes.view(), label, 0);

参数:

res – [输入]
index – [输入]
out_codes – [输出] 目标缓冲区 [n_take, index.pq_dim()]。长度 n_take 定义要解包的记录数，它必须小于列表大小。
label – [输入] 要解码的列表（簇）的 ID。
offset – [输入] 在列表中跳过多少记录。

void unpack_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_vector_view<const uint32_t> in_cluster_indices, raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_codes, uint32_t label )#

根据其列表内偏移量，解压压缩索引中单个列表（簇）的一系列记录，每字节一个编码（与 pq_bits 无关）。

使用示例

// We will unpack the fourth cluster
uint32_t label = 3;
// Create the selection vector
auto selected_indices = raft::make_device_vector<uint32_t>(res, 4);
... fill the indices ...
resource::sync_stream(res);
// allocate the buffer for the output
auto codes = raft::make_device_matrix<float>(res, selected_indices.size(), index.pq_dim());
// decode the whole list
ivf_pq::helpers::codepacker::unpack_list_data(
    res, index, selected_indices.view(), codes.view(), label);

参数:

res – [输入] raft 资源
index – [输入] IVF-PQ 索引（通过引用传递）
in_cluster_indices – [输入] 簇内选定索引的偏移量。
out_codes – [输出] 目标缓冲区 [n_take, index.pq_dim()]。长度 n_take 定义要解包的记录数，它必须小于列表大小。
label – [输入] 要解码的列表（簇）的 ID。

void unpack_contiguous_list_data( raft::resources const &res, const index<int64_t> &index, uint8_t *out_codes, uint32_t n_rows, uint32_t label, uint32_t offset )#

解压压缩索引中从给定 offset 开始的单个列表（簇）的 n_rows 个连续 PQ 编码向量，未扩展到每字节一个编码。输出缓冲区中的每个编码占用 ceildiv(index.pq_dim() * index.pq_bits(), 8) 字节。

使用示例

  raft::resources res;
  // We will unpack the whole fourth cluster
  uint32_t label = 3;
  // Get the list size
  uint32_t list_size = 0;
  raft::update_host(&list_size, index.list_sizes().data_handle() + label, 1,
    raft::resource::get_cuda_stream(res));
  raft::resource::sync_stream(res);
  // allocate the buffer for the output
  auto codes = raft::make_device_matrix<float>(res, list_size, raft::ceildiv(index.pq_dim() *
    index.pq_bits(), 8));
  // unpack the whole list
  ivf_pq::helpers::codepacker::unpack_list_data(res, index, codes.data_handle(), list_size,
label, 0);

参数:

res – [输入] raft 资源
index – [输入] IVF-PQ 索引（通过引用传递）
out_codes – [输出] 目标缓冲区 [n_rows, ceildiv(index.pq_dim() * index.pq_bits(), 8)]。长度 n_rows 定义要解包的记录数，偏移量 + n_rows 必须小于或等于列表大小。
n_rows – [输入] 要解包的编码数
label – [输入] 要解码的列表（簇）的 ID。
offset – [输入] 在列表中跳过多少记录。

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<float, uint32_t, raft::row_major> out_vectors, uint32_t label, uint32_t offset )#

解码压缩索引中从给定 offset 开始的单个列表（簇）的 n_take 个连续记录。

使用示例

  // We will reconstruct the fourth cluster
  uint32_t label = 3;
  // Get the list size
  uint32_t list_size = 0;
  raft::copy(&list_size, index.list_sizes().data_handle() + label, 1,
  resource::get_cuda_stream(res)); resource::sync_stream(res);
  // allocate the buffer for the output
  auto decoded_vectors = raft::make_device_matrix<float>(res, list_size, index.dim());
  // decode the whole list
  ivf_pq::helpers::codepacker::reconstruct_list_data(res, index, decoded_vectors.view(), label,
0);

参数:

res – [输入]
index – [输入]
out_vectors – [输出] 目标缓冲区 [n_take, index.dim()]。长度 n_take 定义要重构的记录数，它必须小于列表大小。
label – [输入] 要解码的列表（簇）的 ID。
offset – [输入] 在列表中跳过多少记录。

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<half, uint32_t, raft::row_major> out_vectors, uint32_t label, uint32_t offset )#

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<int8_t, uint32_t, raft::row_major> out_vectors, uint32_t label, uint32_t offset )#

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_vectors, uint32_t label, uint32_t offset )#

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_vector_view<const uint32_t> in_cluster_indices, raft::device_matrix_view<float, uint32_t, raft::row_major> out_vectors, uint32_t label )#

通过压缩索引中单个列表（聚类）的列表内偏移量解码一系列记录。

使用示例

// We will reconstruct the fourth cluster
uint32_t label = 3;
// Create the selection vector
auto selected_indices = raft::make_device_vector<uint32_t>(res, 4);
... fill the indices ...
resource::sync_stream(res);
// allocate the buffer for the output
auto decoded_vectors = raft::make_device_matrix<float>(
                          res, selected_indices.size(), index.dim());
// decode the whole list
ivf_pq::helpers::codepacker::reconstruct_list_data(
    res, index, selected_indices.view(), decoded_vectors.view(), label);

参数:

res – [输入]
index – [输入]
in_cluster_indices – [输入] 簇内选定索引的偏移量。
out_vectors – [输出] 目标缓冲区 [n_take, index.dim()]。长度 n_take 定义要重构的记录数，它必须小于列表大小。
label – [输入] 要解码的列表（簇）的 ID。

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_vector_view<const uint32_t> in_cluster_indices, raft::device_matrix_view<half, uint32_t, raft::row_major> out_vectors, uint32_t label )#

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_vector_view<const uint32_t> in_cluster_indices, raft::device_matrix_view<int8_t, uint32_t, raft::row_major> out_vectors, uint32_t label )#

void reconstruct_list_data( raft::resources const &res, const index<int64_t> &index, raft::device_vector_view<const uint32_t> in_cluster_indices, raft::device_matrix_view<uint8_t, uint32_t, raft::row_major> out_vectors, uint32_t label )#

void extend_list_with_codes( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_codes, raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices, uint32_t label )#

通过列表标签，在索引中就地扩展一个列表，跳过分类和编码步骤。

使用示例

// We will extend the fourth cluster
uint32_t label = 3;
// We will fill 4 new vectors
uint32_t n_vec = 4;
// Indices of the new vectors
auto indices = raft::make_device_vector<uint32_t>(res, n_vec);
... fill the indices ...
auto new_codes = raft::make_device_matrix<uint8_t, uint32_t, row_major> new_codes(
    res, n_vec, index.pq_dim());
... fill codes ...
// extend list with new codes
ivf_pq::helpers::codepacker::extend_list_with_codes(
    res, &index, codes.view(), indices.view(), label);

参数:

res – [输入]
index – [输入/输出]
new_codes – [输入] 平坦的 PQ 代码，每个字节一个代码 [n_rows, index.pq_dim()]
new_indices – [输入] 源索引 [n_rows]
label – [输入] 目标列表（聚类）的 ID。

void extend_list( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const float, uint32_t, raft::row_major> new_vectors, raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices, uint32_t label )#

通过列表标签，在索引中就地扩展一个列表，跳过分类步骤。

使用示例

// We will extend the fourth cluster
uint32_t label = 3;
// We will extend with 4 new vectors
uint32_t n_vec = 4;
// Indices of the new vectors
auto indices = raft::make_device_vector<uint32_t>(res, n_vec);
... fill the indices ...
auto new_vectors = raft::make_device_matrix<float, uint32_t, row_major> new_codes(
    res, n_vec, index.dim());
... fill vectors ...
// extend list with new vectors
ivf_pq::helpers::codepacker::extend_list(
    res, &index, new_vectors.view(), indices.view(), label);

参数:

res – [输入]
index – [输入/输出]
new_vectors – [输入] 要编码的数据 [n_rows, index.dim()]
new_indices – [输入] 源索引 [n_rows]
label – [输入] 目标列表（聚类）的 ID。

void extend_list( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const int8_t, uint32_t, raft::row_major> new_vectors, raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices, uint32_t label )#

void extend_list( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const uint8_t, uint32_t, raft::row_major> new_vectors, raft::device_vector_view<const int64_t, uint32_t, raft::row_major> new_indices, uint32_t label )#

void erase_list( raft::resources const &res, index<int64_t> *index, uint32_t label )#

从索引中的单个列表（聚类）中移除所有数据。

使用示例

// We will erase the fourth cluster (label = 3)
ivf_pq::helpers::erase_list(res, &index, 3);

参数:

res – [输入]
index – [输入/输出]
label – [输入] 目标列表（聚类）的 ID。

void reset_index(const raft::resources &res, index<int64_t> *index)#

用于重置数据和索引指针以及列表大小的公共辅助 API。在不经过构建阶段的情况下从外部修改索引时非常有用。IVF 列表的数据和索引将丢失。

使用示例

raft::resources res;
using namespace cuvs::neighbors;
// use default index parameters
ivf_pq::index_params index_params;
// initialize an empty index
ivf_pq::index<int64_t> index(res, index_params, D);
// reset the index's state and list sizes
ivf_pq::helpers::reset_index(res, &index);

参数:

res – [输入] raft 资源
index – [输入/输出] 指向 IVF-PQ 索引的指针

void make_rotation_matrix( raft::resources const &res, index<int64_t> *index, bool force_random_rotation )#

公共辅助 API，用于公开计算索引的旋转矩阵。注意：仅当尚未通过 cuvs::neighbors::ivf_pq::build 计算旋转矩阵时使用此 API。

使用示例

raft::resources res;
// use default index parameters
ivf_pq::index_params index_params;
// force random rotation
index_params.force_random_rotation = true;
// initialize an empty index
cuvs::neighbors::ivf_pq::index<int64_t> index(res, index_params, D);
// reset the index
reset_index(res, &index);
// compute the rotation matrix with random_rotation
cuvs::neighbors::ivf_pq::helpers::make_rotation_matrix(
  res, &index, index_params.force_random_rotation);

参数:

res – [输入] raft 资源
index – [输入/输出] 指向 IVF-PQ 索引的指针
force_random_rotation – [输入] 是否对输入数据应用随机旋转矩阵。有关详细信息，请参阅 cuvs::neighbors::ivf_pq::index_params。

void set_centers( raft::resources const &res, index<int64_t> *index, raft::device_matrix_view<const float, uint32_t> cluster_centers )#

用于从外部修改索引的 IVF 中心点的公共辅助 API。注意：在此之前必须重置索引。使用 raft::neighbors::ivf_pq::extend 根据新的中心点构建 IVF 列表。

使用示例

    raft::resources res;
    // allocate the buffer for the input centers
    auto cluster_centers = raft::make_device_matrix<float, uint32_t>(res, index.n_lists(),
index.dim());
    ... prepare ivf centroids in cluster_centers ...
    // reset the index
    reset_index(res, &index);
    // recompute the state of the index
    cuvs::neighbors::ivf_pq::helpers::recompute_internal_state(res, index);
    // Write the IVF centroids
    cuvs::neighbors::ivf_pq::helpers::set_centers(
                   res,
                   &index,
                   cluster_centers);

参数:

res – [输入] raft 资源
index – [输入/输出] 指向 IVF-PQ 索引的指针
cluster_centers – [输入] 新的聚类中心 [index.n_lists(), index.dim()]

void extract_centers( raft::resources const &res, const index<int64_t> &index, raft::device_matrix_view<float, uint32_t, raft::row_major> cluster_centers )#

用于获取已训练索引的 IVF 中心点的公共辅助 API。

使用示例

raft::resources res;
// allocate the buffer for the output centers
auto cluster_centers = raft::make_device_matrix<float, uint32_t>(
  res, index.n_lists(), index.dim());
// Extract the IVF centroids into the buffer
cuvs::neighbors::ivf_pq::helpers::extract_centers(res, index, cluster_centers.data_handle());

参数:

res – [输入] raft 资源
index – [输入] IVF-PQ 索引（通过引用传递）
cluster_centers – [输出] IVF 聚类中心 [index.n_lists(), index.dim]

void recompute_internal_state( const raft::resources &res, index<int64_t> *index )#

辅助方法，用于在 IVF 列表从外部修改后重新计算列表大小及相关数组。

使用示例

using namespace cuvs::neighbors;
raft::resources res;
// use default index parameters
ivf_pq::index_params index_params;
// initialize an empty index
ivf_pq::index<int64_t> index(res, index_params, D);
ivf_pq::helpers::reset_index(res, &index);
// resize the first IVF list to hold 5 records
auto spec = list_spec<uint32_t, int64_t>{
  index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
uint32_t new_size = 5;
ivf::resize_list(res, list, spec, new_size, 0);
raft::update_device(index.list_sizes(), &new_size, 1, stream);
// recompute the internal state of the index
ivf_pq::helpers::recompute_internal_state(res, index);

参数:

res – [输入] raft 资源
index – [输入/输出] 指向 IVF-PQ 索引的指针