类 faiss::gpu::GpuIndexIVFScalarQuantizer
-
class GpuIndexIVFScalarQuantizer : public faiss::gpu::GpuIndexIVF
围绕 GPU 实现的包装器,看起来像 faiss::IndexIVFScalarQuantizer
公共函数
-
GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, const faiss::IndexIVFScalarQuantizer *index, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())
从预先存在的 faiss::IndexIVFScalarQuantizer 实例构造,如果输入索引已训练,则将数据复制到给定的 GPU。
-
GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, int dims, idx_t nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())
构造一个带有空 flat 量化器的新实例;用户提供所需的 IVF 列表的数量。
-
GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, idx_t nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())
构造一个使用提供的 CPU 或 GPU 粗量化器的新实例;用户提供所需的 IVF 列表数量。
-
~GpuIndexIVFScalarQuantizer() override
-
void reserveMemory(size_t numVecs)
在我们的倒排列表中为该数量的向量保留 GPU 内存。
-
void copyFrom(const faiss::IndexIVFScalarQuantizer *index)
从给定的 CPU 索引初始化自身;将覆盖自身中的所有数据
-
void copyTo(faiss::IndexIVFScalarQuantizer *index) const
将自身复制到给定的 CPU 索引;将覆盖索引实例中的所有数据
-
size_t reclaimMemory()
添加向量后,可以调用此函数来回收设备内存,使其精确到所需的量。返回以字节为单位回收的空间
-
virtual void reset() override
清除所有倒排列表,但保留粗量化器和标量量化器信息
-
virtual void updateQuantizer() override
如果用户手动更改 IVF 粗量化器的状态(例如,替换新实例或在训练范围之外更改粗量化器中的向量),则应调用此函数
-
virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const
返回特定倒排列表中包含的编码向量数据,用于调试目的。如果 gpuFormat 为 true,则返回的数据按 GPU 端表示形式进行编码。否则,它将转换为 CPU 格式。兼容格式,而本机 GPU 格式可能不同。
-
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
查询维度为 d 的 n 个向量到索引。
返回所有距离 < 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)
将条目从另一个数据集移动到自身。输出时,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 所需的字节数
公共成员
-
faiss::ScalarQuantizer sq
与 CPU 版本类似。
-
bool by_residual
与 CPU 版本类似。
-
int d
向量维度。
-
bool verbose
详细程度级别。
-
MetricType metric_type
此索引用于搜索的度量类型。
-
float metric_arg
度量类型的参数。
-
size_t nprobe = 1
查询时的探针数量。
-
size_t max_codes = 0
查询时要访问的最大代码数。
-
size_t nlist = 0
倒排列表的数量。
-
char quantizer_trains_alone = 0
>= 0:在 kmeans 训练中使用量化器作为索引 = 1:只需将训练集传递给量化器的 train() = 2:在平面索引上进行 kmeans 训练 + 将质心添加到量化器。
-
bool own_fields = false
对象是否拥有量化器。
-
ClusteringParameters cp
用于覆盖默认的聚类参数。
受保护的函数
-
void verifySQSettings_() const
验证索引 SQ 参数。
-
int getCurrentNProbe_(const SearchParameters *params) const
从当前设置的 nprobe 或 SearchParameters(如果可用)返回应该用于当前搜索的 nprobe。
-
void verifyIVFSettings_() const
-
virtual bool addImplRequiresIDs_() const override
addImpl_ 是否需要 IDs? 如果是,并且没有提供 IDs,我们将根据添加 IDs 的顺序依次生成它们。
-
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 GpuIndexIVFScalarQuantizerConfig ivfSQConfig_
我们的配置选项。
-
size_t reserveMemoryVecs_
期望的倒排列表内存预留。
-
const GpuIndexIVFConfig ivfConfig_
我们的配置选项。
-
std::shared_ptr<GpuResources> resources_
管理设备的流、cuBLAS 句柄和暂存内存。
-
const GpuIndexConfig config_
我们的配置选项。
-
size_t minPagedSize_
从此大小开始,我们将CPU到GPU的复制进行分页。
-
GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, const faiss::IndexIVFScalarQuantizer *index, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())