类 faiss::gpu::GpuIndexIVFPQ

class GpuIndexIVFPQ : public faiss::gpu::GpuIndexIVF 

GPU 的 IVFPQ 索引。

公共类型

using component_t = float

using distance_t = float

公共函数

GpuIndexIVFPQ(GpuResourcesProvider *provider, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 从预先存在的 faiss::IndexIVFPQ 实例构造，如果输入索引已训练，则将数据复制到给定的 GPU。

GpuIndexIVFPQ(GpuResourcesProvider *provider, int dims, idx_t nlist, idx_t subQuantizers, idx_t bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 构造具有空 flat 量化器的新的实例；用户提供所需的 IVF 列表的数量。

GpuIndexIVFPQ(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, idx_t nlist, idx_t subQuantizers, idx_t bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 构造具有提供的 CPU 或 GPU 粗量化器的新的实例；用户提供所需的 IVF 列表的数量。

~GpuIndexIVFPQ() override

void copyFrom(const faiss::IndexIVFPQ *index): 在 GPU 上为倒排列表预留空间，用于 num 个向量，假设这些向量均匀分布。从给定的 CPU 索引初始化自身；将覆盖自身中的所有数据。

void copyTo(faiss::IndexIVFPQ *index) const: 将自身复制到给定的 CPU 索引；将覆盖索引实例中的所有数据。

void reserveMemory(size_t numVecs): 在我们的倒排列表中为该数量的向量预留 GPU 内存。

void setPrecomputedCodes(bool enable): 启用或禁用预计算代码。

bool getPrecomputedCodes() const: 是否启用了预计算代码？

int getNumSubQuantizers() const: 返回我们正在使用的子量化器的数量。

int getBitsPerCode() const: 返回每个 PQ 代码的位数。

int getCentroidsPerSubQuantizer() const: 返回每个 PQ 代码的质心数量（2^每个代码的位数）。

size_t reclaimMemory(): 添加向量后，可以调用此函数来回收设备内存，使其精确到所需的量。返回以字节为单位回收的空间。

virtual void reset() override: 清除所有倒排列表，但保留粗糙和乘积质心信息。

virtual void updateQuantizer() override: 如果用户手动更改 IVF 粗糙量化器的状态（例如，在训练范围之外替换新的实例或更改粗糙量化器中的向量），则应调用此函数。

virtual void train(idx_t n, const float *x) override: 基于给定的向量数据训练粗糙量化器和乘积量化器。

void copyFrom(const faiss::IndexIVF *index): 从 CPU 等效对象复制我们需要的内容。

void copyTo(faiss::IndexIVF *index) const: 将我们拥有的内容复制到 CPU 等效对象。

virtual idx_t getNumLists() const: 返回我们正在管理的倒排列表的数量。

virtual idx_t getListLength(idx_t listId) const: 返回特定倒排列表中存在的向量的数量。

virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const: 返回特定倒排列表中包含的编码向量数据，用于调试目的。如果 gpuFormat 为 true，则数据将以 GPU 端表示中编码的形式返回。否则，它将被转换为 CPU 格式。兼容格式，而本地 GPU 格式可能不同。

virtual std::vector<idx_t> getListIndices(idx_t listId) const: 返回特定倒排列表中包含的向量索引，用于调试目的。

virtual void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr, IndexIVFStats *stats = nullptr) const override

搜索一组向量，这些向量已由 IVF 量化器预量化。用查询结果填充相应的堆。默认实现使用 InvertedListScanners 进行搜索。

参数:

n – 要查询的向量数
x – 查询向量，大小为 nx * d
assign – 粗量化索引，大小为 nx * nprobe
centroid_dis – 到粗质心的距离，大小为 nx * nprobe
distance – 输出距离，大小为 n * k
labels – 输出标签，大小为 n * k
store_pairs – 在结果的上/下 32 位中存储 inv list index + inv list offset，而不是 ids（用于重新排序）。
params – 用于覆盖对象的搜索参数
stats – 要更新的搜索统计信息（可以为 null）

virtual void range_search_preassigned(idx_t nx, const float *x, float radius, const idx_t *keys, const float *coarse_dis, RangeSearchResult *result, bool store_pairs = false, const IVFSearchParameters *params = nullptr, IndexIVFStats *stats = nullptr) const override

范围搜索一组向量，这些向量已由 IVF 量化器预量化。填充 RangeSearchResults 结果。默认实现使用 InvertedListScanners 进行搜索。

参数:

n – 要查询的向量数
x – 查询向量，大小为 nx * d
assign – 粗量化索引，大小为 nx * nprobe
centroid_dis – 到粗质心的距离，大小为 nx * nprobe
result – 输出结果
store_pairs – 在结果的上/下 32 位中存储 inv list index + inv list offset，而不是 ids（用于重新排序）。
params – 用于覆盖对象的搜索参数
stats – 要更新的搜索统计信息（可以为 null）

int getDevice() const: 返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources(): 返回对我们的 GpuResources 对象的引用，该对象管理 GPU 上的内存、流和句柄资源。

void setMinPagingSize(size_t size): 设置我们使用 CPU -> GPU 分页进行搜索的最小数据大小（以 MiB 为单位）

size_t getMinPagingSize() const: 返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override: x 可以驻留在 CPU 或任何 GPU 上；根据需要执行复制。如果添加集太大，则处理分页添加；调用 addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override: x 和 ids 可以驻留在 CPU 或任何 GPU 上；根据需要执行复制。如果添加集太大，则处理分页添加；调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override: x 和 labels 可以驻留在 CPU 或任何 GPU 上；根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override: x、distances 和 labels 可以驻留在 CPU 或任何 GPU 上；根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override: x、distances 和 labels 以及 recons 可以驻留在 CPU 或任何 GPU 上；根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override: 重写以强制 GPU 索引提供其自己的 GPU 友好的实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override: 重写以强制 GPU 索引提供其自己的 GPU 友好的实现

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

查询 n 个维度为 d 的向量到索引。

返回所有距离 < radius 的向量。请注意，许多索引未实现 range_search（只有 k-NN 搜索是强制性的）。

参数:

n – 向量的数量
x – 要搜索的输入向量，大小为 n * d
radius – 搜索半径
result – 结果表

virtual size_t remove_ids(const IDSelector &sel): 从索引中删除 ID。并非所有索引都支持。返回删除的元素数量。

virtual void reconstruct(idx_t key, float *recons) const

重建存储的向量（如果是有损编码，则为近似值）

对于某些索引，可能未定义此函数

参数:

key – 要重建的向量的 ID
recons – 重建的向量（大小为 d）

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

重建多个存储的向量（如果是有损编码，则为近似值）

对于某些索引，可能未定义此函数

参数:

n – 要重建的向量的数量
keys – 要重建的向量的 ID（大小为 n）
recons – 重建的向量（大小为 n * d）

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

重建向量 i0 到 i0 + ni - 1

对于某些索引，可能未定义此函数

参数:

i0 – 序列中第一个向量的索引
ni – 序列中向量的数量
recons – 重建的向量（大小为 ni * d）

virtual DistanceComputer *get_distance_computer() const

为此类索引获取 DistanceComputer（在 AuxIndexStructures 中定义）对象。

对于支持对其向量进行随机访问的索引，实现了 DistanceComputer。

virtual size_t sa_code_size() const: 生成的代码的大小（以字节为单位）

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:

n – 向量的数量
x – 输入向量，大小为 n * d
bytes – 输出编码的向量，大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:

n – 向量的数量
bytes – 输入编码向量，大小为 n * sa_code_size()
x – 输出向量，大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0): 将条目从另一个数据集移动到 self。输出时，other 为空。 add_id 被添加到所有移动的 id (对于顺序 id，这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const: 检查两个索引是否兼容（即，它们以相同的方式训练并且具有相同的参数）。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:

codes – 要添加的代码，大小为 n * sa_code_size()
xids – 相应的 id，大小为 n

void train_q1(size_t n, const float *x, bool verbose, MetricType metric_type): 训练量化器并调用 train_residual 来训练子量化器。

size_t coarse_code_size() const: 计算存储列表 ID 所需的字节数

void encode_listno(idx_t list_no, uint8_t *code) const

idx_t decode_listno(const uint8_t *code) const

公共成员

ProductQuantizer pq: 像 CPU 版本一样，我们公开一个公开可见的 ProductQuantizer 用于操作

int d: 向量维度

idx_t ntotal: 索引向量的总数

bool verbose: 详细级别

bool is_trained: 如果 Index 不需要训练，或者如果训练已经完成，则设置

MetricType metric_type: 此索引用于搜索的度量类型

float metric_arg: 度量类型的参数

size_t nprobe = 1: 查询时的 probes 数量

size_t max_codes = 0: 查询时访问的最大代码数

Index *quantizer = nullptr: 将向量映射到倒排列表的量化器

size_t nlist = 0: 倒排列表的数量

char quantizer_trains_alone = 0: = 0: 将量化器用作 kmeans 训练中的索引 = 1: 仅将训练集传递给量化器的 train() = 2: 在平面索引上进行 kmeans 训练 + 将质心添加到量化器

bool own_fields = false: 对象是否拥有量化器

ClusteringParameters cp: 覆盖默认聚类参数

Index *clustering_index = nullptr: 覆盖聚类期间使用的索引

受保护的函数

void setIndex_(GpuResources *resources, int dim, idx_t nlist, faiss::MetricType metric, float metricArg, int numSubQuantizers, int bitsPerSubQuantizer, bool useFloat16LookupTables, bool useMMCodeDistance, bool interleavedLayout, float *pqCentroidData, IndicesOptions indicesOptions, MemorySpace space): 初始化适当的索引。

void verifyPQSettings_() const: 如果配置设置不正确，则抛出错误。

void trainResidualQuantizer_(idx_t n, const float *x): 根据给定的向量数据训练 PQ 量化器。

void copyFrom(const faiss::Index *index): 从 CPU 等效对象复制我们需要的内容。

void copyTo(faiss::Index *index) const: 将我们拥有的内容复制到 CPU 等效对象。

int getCurrentNProbe_(const SearchParameters *params) const: 从当前集合 nprobe 或 SearchParameters（如果可用）中返回我们应该用于当前搜索的 nprobe

void verifyIVFSettings_() const

virtual bool addImplRequiresIDs_() const override: addImpl_ 是否需要 ID？如果是，且未提供 ID，我们将根据添加 ID 的顺序依次生成它们

virtual void trainQuantizer_(idx_t n, const float *x)

virtual void addImpl_(idx_t n, const float *x, const idx_t *ids) override: 从 GpuIndex 中调用以进行 add/add_with_ids。

virtual void searchImpl_(idx_t n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override: 从 GpuIndex 中调用以进行搜索。

受保护的属性

const GpuIndexIVFPQConfig ivfpqConfig_: 我们在初始化时使用的配置选项。

bool usePrecomputedTables_: 运行时覆盖：是否使用预计算表。

int subQuantizers_: 每个编码向量的子量化器数量。

int bitsPerCode_: 每个子量化器代码的位数。

size_t reserveMemoryVecs_: 所需的倒排列表内存预留。

std::shared_ptr<IVFPQ> index_: 我们拥有的乘积量化器实例；包含倒排列表

const GpuIndexIVFConfig ivfConfig_: 我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_: 对于经过训练/初始化的索引，这是对基类的引用。

std::shared_ptr<GpuResources> resources_: 管理设备的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_: 我们的配置选项。

size_t minPagedSize_: 大于此大小，我们将从 CPU 分页复制到 GPU。