4. mcFaiss编程指南

4.1. mcFaiss索引接口

使用mcFaiss时，通常遵循的接口调用顺序为：创建索引 -> 训练索引 -> 添加向量数据 -> 搜索向量。

4.2. 索引的基本使用

本节介绍如何在曦云系列GPU上使用mcFaiss GPU索引。

4.2.1. 创建索引

创建一个索引类。

代码示例（C++）

使用C++ API时，首先需要创建一个Index类的实例：

int d = 4; // dimension
int nlist = 1024; // number of clusters

faiss::gpu::StandardGpuResources res;
faiss::gpu::GpuIndexIVFFlatConfig config;
faiss::gpu::GpuIndexIVFFlat index(&res, d, nlist, faiss::METRIC_L2, config);

示例中，创建了一个GPU索引类，为使用欧式距离（L2）的IVFFlat索引类，nlist为1024。关于nlist的具体含义，参见 5.2 GpuIndexIVFFlat。

代码示例（Python）

d = 4 # dimension
nlist = 1024 # number of clusters

config = faiss.GpuIndexIVFFlatConfig()
res = faiss.StandardGpuResources()
index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, config)

4.2.2. 训练索引

对于某些索引，必须对其进行训练。比如，对于IVF索引，进行训练的目的是为了生成索引的聚类中心，从而提高索引的搜索效率。

代码示例（C++）

int nb = 10000; // 训练数据的数量
std::vector<float> xb(d * nb);
// 填充训练的向量数据到xb
index.train(nb, xb.data());

代码示例（Python）

nb = 10000 # 训练数据的数量
# 填充训练的向量数据到xb，这里使用随机数做示例
xb = np.random.random((nb, d)).astype('float32')
index.train(xb)

在实际使用过程中，并不要求训练的数据和 4.2.3 添加向量数据中添加的数据是一致的，只要保证这两者的维度是一致的即可。 但是为了提高索引的搜索效率，建议训练的数据和添加的数据一致，或者训练的数据是添加数据的一个子集。

对于Flat索引，不需要进行训练，因为Flat索引不需要聚类中心。

4.2.3. 添加向量数据

将向量数据添加到索引中。

代码示例（C++）

int nb = 10000; // 向量数据的数量
std::vector<float> xb(d * nb);
// 填充向量数据到xb
index.add(nb, xb.data());

代码示例（Python）

nb = 10000 # 向量数据的数量
# 填充向量数据到xb，这里使用随机数做示例
xb = np.random.random((nb, d)).astype('float32')
index.add(xb)

添加操作可以反复进行，但是每次添加都有可能触发设备内存的重新分配。因此，为了提高数据添加的效率，建议批量添加数据。

4.2.4. 搜索向量

向量搜索是索引的核心功能，可以通过以下方式进行搜索。

代码示例（C++）

int nq = 10; // 查询向量的数量
int topk = 3; // 搜索的topk
std::vector<float> xq(d * nq);
// 填充查询向量数据到xq
std::vector<float> dis(d * nq * topk);
std::vector<int64_t> rids(d * nq * topk);

index.search(nq, xq.data(), topk, dis.data(), rids.data());

搜索结束后，dis和rids中保存了搜索结果，dis中保存了搜索结果的距离，rids中保存了搜索结果的ID。

代码示例（Python）

nq = 10 # 查询向量的数量
topk = 3 # 搜索的topk
# 填充查询向量数据到xq，这里使用随机数做示例
xq = np.random.random((nq, d)).astype('float32')

dis, rids = index.search(xq, topk)

通常情况下，GPU上的搜索效率会比CPU上的更高，但是在搜索之前需要将查询向量数据拷贝到GPU上，因此在实际使用的过程中，为了提高搜索的效率，建议进行批量搜索。