Faiss(3)：基于IndexIVFPQ的demo程序

1. 说明

在研究Faiss库一段时间之后，做了一个简单的demo程序。这个程序在CPU环境下进行训练，GPU内进行搜索。使用IndexIVFPQ索引。

由于没有实际的图片库，这里使用numpy生成多维随机数组的方式来模拟数据检索的过程。

1.1 开发环境

参考Faiss(2)：编程环境搭建

1.2 程序总体流程

Faiss总体使用过程可以分为三步：

1. 构建训练数据和查询数据，原始数据集以二维矩阵形式表达,矩阵每一行代表数据库中的一个数据，每一列代表数据库中所有数据的某个特征，每个数据为高维浮点矢量如64维，当数据库中数据量较大时，训练集规模远小于数据库。查询数据与训练数据类型和格式相同。这一步骤是离线进行，更新频率很低。

2. 挑选合适配置参数建立Index（Faiss的核心部件），从数据库中选取训练数据训练并得到聚类中心向量，数据库数据通过add操作建立Index索引。这部分也是离线，更新频率高于上一步。

3. Search搜索，距离计算+比较得到最后结果。在线进行，目标是低延迟，高吞吐量。

2. 程序设计

1. 设置参数配置

nb = nb            # dataset numbers
nq = nq            # query numbers
nlist = nlist
nprobe = nprobe
topk = topk
d = 64             # dimension
m = 64             # sub quantizer
nbits = 8          # bits of per code

learning_ratio = lr   # 学习率

程序使用固定为64的维度， 其中nb表示database，本文以1M数据量为例进行说明。

2. 准备测试样本集和训练集

# prep nq test samples
def CreateDataSet(nb, nq, d):
    # dataset
    np.set_printoptions(threshold = np.inf)
    np.random.seed(1234)

    xb = np.random.random((nb, d)).astype('float32')
    xb[:, 0] += np.arange(nb) / int(nb)
    
    # query
    xq = np.random.random((nq, d)).astype('float32')
    # xq[:, 0] += np.random.randint(low=0, high=1e6, size=nq) / int(1e5)
    xq[:, 0] += np.arange(nq) * 1.0/nq
    
    # prep learning
    learning = int(nb*learning_ratio)
    np.random.seed(900)
    learning_d = np.random.random((learning,d)).astype('float32')
    learning_d[:,0] += np.arange(learning) * 1.0/learning

    return xb, xq， learning_d

这里nq表示要查询的目标向量集中向量的个数。后续查询时，以向量集xq目标向量进行查询。
learning_d是用于训练的训练集，总数是database的十分之一 (nb * lr)，但是由不同的random seed产生不同的数据。

3. 创建index实例

def CreateIndex(nlist, d, m, nbits):
    quantizer = faiss.IndexFlatL2(d)
    index = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)

    return index

IndexFlatL2本身也是一个索引，它是最简单的索引类型，只执行强力L2距离搜索。但是这里为了扩展到非常大的数据集，将其作为一个量化器来压缩存储的向量的变体。压缩的方法基于乘积量化。损失一定精度为代价， 自身距离也不为0， 这是由于有损压缩。

4. 训练、添加向量和搜索

# training
index.train(learning_d)

# add dataset
index.add(xb)

# search
D, I = index.search(xq, topk)

我这里使用的训练集是数据集的1/10。
search函数返回两个列表，D表示search结果的距离（float型）， I表示search结果向量的id号(int型)。

search函数（C++源代码中的实现）：

virtual void search (idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const = 0;
/*
 * n : 要查询的向量个数
 * x : 输入向量集，即要检索的向量，size: n*d
 * k : 输出向量集个数，即k邻
 * distances : 输出向量集对应的距离向量，size: n*k
 * labels : 输出向量集，size: n*k
*/

在python的接口中只留出两个参数：x和k，也就是说上述python代码分1000次检索xq_t[x]，输出100个近邻结果。

5.拷贝index到GPU中

co = faiss.GpuClonerOptions()
co.useFloat16 = True
res = faiss.StandardGpuResources()
index_gpu = faiss.index_cpu_to_gpu(res, 0, index, co)

这一步其实在创建index实例之后就可以进行，那么后续的train、add则直接调用index_gpu在GPU中进行，也可以在cpu中add和train之后再整体拷贝到GPU中。

由于我用的Tesla P4的shared memory只有48M，而当m设置为64时所需的shared memory为64M，所以我通过co.useFloat16=True来“useFloat16LookupTables”，否则会报如下错误：

Traceback (most recent call last):
  File "faiss/GPU_tim.py", line 132, in <module>
    main(nb, seg, lr)
  File "faiss/GPU_time.py", line 94, in main
    index
  File "/home/montage/.local/lib/python2.7/site-packages/faiss/__init__.py", line 485, in index_cpu_to_all_gpus
    index2 = index_cpu_to_gpu_multiple_py(res, index, co)
  File "/home/montage/.local/lib/python2.7/site-packages/faiss/__init__.py", line 477, in index_cpu_to_gpu_multiple_py
    index = index_cpu_to_gpu_multiple(vres, vdev, index, co)
RuntimeError: Error in void faiss::gpu::GpuIndexIVFPQ::verifySettings_() const at gpu/GpuIndexIVFPQ.cu:432: Error: 'requiredSmemSize <= getMaxSharedMemPerBlock(device_)' failed: Device 0 has 49152 bytes of shared memory, while 8 bits per code and 64 sub-quantizers requires 65536 bytes. Consider useFloat16LookupTables and/or reduce parameters

当然，也可以通过减小m的值来避免此项错误。

6. 计算结果的精确度

IndexIVFPQ的搜索结果的精确度是通过将结果与暴力搜索的结果进行比较得出的，将两者的结果求交集，可以知道每次搜索中有多少是重复的结果，再将交集数量除以topk就是精确度了。

暴力搜索可以直接使用IndexFlatL2索引进行，由于暴力搜索是简单的对向量进行依次比较，所以可以不进行训练。

# 暴力搜索
def ViolenceSearch(xb, xq, d, topk):
    print("start violence search")
    quantizer = faiss.IndexFlatL2(d)
    quantizer.add(xb)
    D, I = quantizer.search(xq, topk)

def CalAccuracy(nq, I, refI, topk):
    print("calculate accuracy...")
    rec = np.zeros(nq)
    for i in range(nq):
        rec[i] = 1.0 * len(set(I[i][:]).intersection(set(refI[i][:])))/topk
    
    accu = sum(rec)*1./nq * 100
    return accu

注：为了在多种参数配置下节省时间，我将各个步骤进行单独封装，以便调用，在实际开发过程中可根据需要，不一定需要按照相同步骤进行。

3. 运行程序及结果

略