类 faiss::gpu::GpuIndex

class GpuIndex : public faiss::Index

子类:faiss::gpu::GpuIndexCagra, faiss::gpu::GpuIndexFlat, faiss::gpu::GpuIndexIVF

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndex(std::shared_ptr<GpuResources> resources, int dims, faiss::MetricType metric, float metricArg, GpuIndexConfig config)
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

xids 可以在 CPU 或任何 GPU 上;根据需要执行复制 如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 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

xdistanceslabels 可以驻留在 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

xdistanceslabelsrecons 可以驻留在 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 train(idx_t n, const float *x)

对一组有代表性的向量执行训练

参数:
  • n – 训练向量的数量

  • x – 训练向量,大小为 n * d

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 void reset() = 0

从数据库中移除所有元素。

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

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细程度

bool is_trained

如果 Index 不需要训练,或者已经完成训练,则设置此标志

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

保护函数

void copyFrom(const faiss::Index *index)

从 CPU 等效项复制我们需要的内容。

void copyTo(faiss::Index *index) const

将我们拥有的内容复制到 CPU 等效项。

virtual bool addImplRequiresIDs_() const = 0

addImpl_ 是否需要 ID?如果是,且未提供 ID,我们将根据添加 ID 的顺序依次生成它们

virtual void addImpl_(idx_t n, const float *x, const idx_t *ids) = 0

被覆盖以实际执行添加 所有数据都保证驻留在我们的设备上

virtual void searchImpl_(idx_t n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const = 0

被覆盖以实际执行搜索 所有数据都保证驻留在我们的设备上

保护属性

std::shared_ptr<GpuResources> resources_

管理设备的流、cuBLAS 句柄和临时内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

大于此大小的数据,我们将从 CPU 页面式复制到 GPU。

私有函数

void addPaged_(idx_t n, const float *x, const idx_t *ids)

如果添加集太大,则处理页面式添加,并传递给 addImpl_ 以实际执行当前页面的添加

void addPage_(idx_t n, const float *x, const idx_t *ids)

为 GPU 驻留数据的单个页面调用 addImpl_。

void searchNonPaged_(idx_t n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

为 GPU 驻留数据的单个页面调用 searchImpl_。

void searchFromCpuPaged_(idx_t n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

为 GPU 驻留数据的单个页面调用 searchImpl_,处理数据的分页和从 CPU 的复制