Faiss学习：一

在多个GPU上运行Faiss以及性能测试

一、Faiss的基本使用

1.1在CPU上运行

Faiss的所有算法都是围绕index展开的。不管运行搜索还是聚类，首先都要建立一个index。

import faiss
# make faiss available
index = faiss.IndexFlatL2(d)
# build the index
# d is the dimension of data

在运行上述代码后，就可以添加数据并运行搜索了。

index.add(xb)
# xb is the base data
D, I = index.search(xq, k)
# xq is the query data
# k is the num of neigbors you want to search
# D is the distance matrix between xq and k neigbors
# I is the index matrix of k neigbors

1.2在单个GPU上运行

在单个GPU上运行的语法基本与在GPU上运行类似。但是需要申明一个GPU资源的标识.

res = faiss.StandardGpuResources()
# we need only a StandardGpuResources per GPU
flat_config = 0
# flat_config is an ID. if you have 3 GPUs, flat_configs maybe 0, 1, 2
index = faiss.GpuIndexFlatL2(res, d, flat_config)
# build the index
index.add(xb)
D, I = index.search(xq, k)

1.3在多个GPU上运行

在多个GPU上运行时便有所不同，我们需要将数据集分割给多个GPU以完成并行搜索。

在Faiss中提供了两种方法实现：IndexProxy和IndexShards。

下面着重介绍IndexProxy。

res = [faiss.StandardGpuResources() for i in range(ngpu)]
# first we get StandardGpuResources of each GPU
# ngpu is the num of GPUs

indexes = [faiss.GpuIndexFlatL2(res[i], i, d, useFloat16)
           for i in range(ngpu)]
# then we make an Index array
# useFloat16 is a boolean value

index = faiss.IndexProxy()
for sub_index in indexes:
    index.addIndex(sub_index)
# build the index by IndexProxy