GPU与CPU版Faiss索引检索耗时相同的原因及代码排查

我有一个包含500万样本的数据集，需要为每个样本检索数据集中最相似的20个样本。为测试效率，我分别编写了GPU版与CPU版Faiss索引代码，但在V100机器上运行时，两段代码均耗时约25分钟。想知道为什么GPU与CPU版索引的查询耗时一致，是不是代码存在问题？

GPU版代码 #

import numpy as np
import faiss
import time
from tqdm import tqdm
import math

def geshihua(total_sec):
    h=total_sec//3600
    total_sec-=h*3600
    minutes = total_sec//60
    total_sec-=minutes*60
    return h,minutes,total_sec

start = int(time.time())
start_time = time.strftime('%Y-%m-%d %H:%M:%s',time.localtime())
print("start on",start_time)

# produce 5000000 vectors
total_samples_num=5000000
dataset = np.random.random((total_samples_num, 128)).astype('float32')

# definite faiss index
res = faiss.StandardGpuResources()
dim,measure=128,faiss.METRIC_L2
param = 'IVF100,PQ16'
index = faiss.index_factory(dim,param,measure)
gpu_index = faiss.index_cpu_to_gpu(res,0,index)
print(gpu_index.is_trained)
# train the faiss
gpu_index.train(dataset)
# add data to dataset
gpu_index.add(dataset)


retri_num = 20
total_indices=np.random.rand(0,retri_num)
bs=1000
for i in tqdm(range(int(math.ceil(total_samples_num/bs)))):
# for i in tqdm(range(int(total_samples_num/bs))):
    queryset=dataset[i*bs:min((i+1)*bs,total_samples_num-i*bs)]
    # use faiss index to get 20 the most similar vector
    distances, indices = gpu_index.search(queryset, retri_num)
    total_indices=np.vstack((total_indices,indices))
    end_time = time.strftime('%Y-%m-%d %H:%M:%s',time.localtime())

end=int(time.time())
total_sec=end-start
h,minutes,total_sec=geshihua(total_sec)
print(f"total running time:{h}h {minutes}m {total_sec}s")

CPU版代码 #

import numpy as np
import faiss
import time
from tqdm import tqdm
import math

def geshihua(total_sec):
    h=total_sec//3600
    total_sec-=h*3600
    minutes = total_sec//60
    total_sec-=minutes*60
    return h,minutes,total_sec

start = int(time.time())
start_time = time.strftime('%Y-%m-%d %H:%M:%s',time.localtime())
print("start on",start_time)

# produce 5000000 vectors
total_samples_num=5000000
dataset = np.random.random((total_samples_num, 128)).astype('float32')

# definite faiss index
res = faiss.StandardGpuResources()
dim,measure=128,faiss.METRIC_L2
param = 'IVF100,PQ16'
index = faiss.index_factory(dim,param,measure)
# index = faiss.index_cpu_to_gpu(res,0,index)
print(index.is_trained)
# train the faiss
index.train(dataset)
# add data to dataset
index.add(dataset)

retri_num = 20
total_indices=np.random.rand(0,retri_num)
bs=1000
for i in tqdm(range(int(math.ceil(total_samples_num/bs)))):
# for i in tqdm(range(int(total_samples_num/bs))):
    queryset=dataset[i*bs:min((i+1)*bs,total_samples_num-i*bs)]
    # use faiss index to get 20 the most similar vector
    distances, indices = index.search(queryset, retri_num)
    total_indices=np.vstack((total_indices,indices))
    end_time = time.strftime('%Y-%m-%d %H:%M:%s',time.localtime())
    # print(f"finish {i*bs}~{(i+1)*bs} on {end_time}")

end=int(time.time())
total_sec=end-start
h,minutes,total_sec=geshihua(total_sec)
print(f"total running time:{h}h {minutes}m {total_sec}s")

原因分析与优化建议 #

核心原因 #

1. 索引参数未适配GPU特性
   你使用的IVF100,PQ16中，IVF聚类中心仅100个，对于500万样本来说，每个聚类下平均有5万个样本。GPU的优势是大规模并行计算，过小的聚类数会导致每个查询需要遍历大量聚类内样本，并行能力无法充分发挥，最终GPU的计算优势被抵消。

2. 内存操作与数据传输成为瓶颈

   • GPU版中，每次查询的queryset从CPU内存切片后传输到GPU，频繁小批量数据传输带来额外开销。
   • total_indices=np.vstack((total_indices,indices))在CPU上执行，每次都要重新分配内存并复制数据，这部分耗时可能远超过GPU查询本身，导致整体耗时和CPU版持平。

3. CPU版已利用多线程并行
   Faiss的CPU索引默认开启多线程计算，V100机器通常搭配高性能多核CPU，CPU版已经利用所有核心并行处理，所以耗时与GPU版接近。

优化建议 #

1. 调整索引参数
   增大IVF聚类数量，比如改为IVF2048,PQ16或IVF4096,PQ16，减少每个聚类的样本数，让GPU的并行计算能力得到充分发挥。

2. 优化内存与数据传输

   • 预先将整个dataset复制到GPU内存，避免每次查询的数据传输：

     dataset_gpu = faiss.index_cpu_to_gpu(res, 0, dataset)
     

   • 预先分配total_indices内存，用切片赋值替代np.vstack：

     total_indices = np.empty((total_samples_num, retri_num), dtype=np.int64)
     # 循环内替换为
     total_indices[i*bs:min((i+1)*bs, total_samples_num)] = indices
     

3. 优化GPU资源配置
   为GPU索引配置更多临时显存并启用半精度计算，提升性能：

   res = faiss.StandardGpuResources()
   res.setTempMemory(2 * 1024 * 1024 * 1024)  # 分配2GB临时显存
   config = faiss.GpuIndexIVFPQConfig()
   config.useFloat16 = True  # 启用半精度计算
   gpu_index = faiss.index_cpu_to_gpu(res, 0, index, config)
   

4. 增大查询批次
   将batch size从1000调整为8192或16384，减少循环次数和数据传输的开销。

内容的提问来源于stack exchange，提问作者袁梓晽