在多個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
在運行上述代碼后鹃两,就可以添加數(shù)據(jù)并運行搜索了。
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資源的標(biāo)識.
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上運行時便有所不同扼仲, 我們需要將數(shù)據(jù)集分割給多個GPU以完成并行搜索。
在Faiss中提供了兩種方法實現(xiàn):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
二、kmeans測試
如圖所示數(shù)據(jù)為1M個肆资,中心點為1K個矗愧。
在不同數(shù)據(jù)維度以及GPU數(shù)目下迭代20次所需要的時間。
三郑原、暴力搜索測試
數(shù)據(jù)集為sift1M唉韭, 該數(shù)據(jù)集共1M個,128維犯犁。(運行在兩個K40M GPU上)
可以看到在每次查詢10K個數(shù)據(jù)的1024個最近鄰居時平均每個查詢只需360ns属愤。當(dāng)需要查詢的鄰居數(shù)下降時,查詢時間能夠降至100ns酸役。
四住诸、IVFPQ搜索測試
數(shù)據(jù)集同上,運行環(huán)境同上涣澡。
基本參數(shù):
numCentroids=4096
numQuantizers=64
首先我們測試nprob對性能的影響
當(dāng)nprob上升時贱呐, 每次查詢時間會增加, 同時查詢的準(zhǔn)確度也會上升入桂。但上升到一定程度上升幅度便會迅速變小吼句。我們?nèi)?zhǔn)確度的拐點值nprob=32進行下一步測試。
接下來我們測試查詢的鄰居數(shù)即k值對性能的影響事格。
可以看到查詢時間不再是線性增長了惕艳。也就意味著對于IVFPQ鄰居數(shù)不宜太多。
選定k=32進行下一步測試驹愚。
如圖所示远搪,隨著每次查詢的數(shù)量上升,平均查詢時間先變小在變大逢捺,這可能是由于數(shù)據(jù)量小時開銷比較大導(dǎo)致平均查詢時間較大谁鳍。可以看到隨著查詢的數(shù)量上升劫瞳, 平均查詢時間上升但上升幅度放緩倘潜,估計會在250ns左右穩(wěn)定。