这个脚本将测量使用GPU进行搜索的时间。如果搜索时间显著减少,那么GPU版本的Faiss应该正常工作。
本人实测,gpu版本的输出应该如下:
Faiss支持GPU
搜索时间: 0.0423 秒
cpu的输出为:
Faiss不支持GPU
搜索时间: 66.4641 秒
python
import faiss
import numpy as np
import time
d = 64 # 向量维度
nb = 100000 # 数据集大小
nq = 10000 # 查询集大小
np.random.seed(0) # 设置随机种子
# 检查GPU支持
if faiss.get_num_gpus() > 0:
print("Faiss支持GPU")
else:
print("Faiss不支持GPU")
# 创建数据集和查询集
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
# 使用GPU进行搜索
index = faiss.IndexFlatL2(d) # 使用L2距离
if faiss.get_num_gpus() > 0:
res = faiss.StandardGpuResources() # 使用默认的GPU资源
index = faiss.index_cpu_to_gpu(res, 0, index)
# 计时搜索
start_time = time.time()
index.add(xb)
D, I = index.search(xq, 10) # 搜索最近的10个邻居
end_time = time.time()
# 打印搜索时间
print(f"搜索时间: {end_time - start_time:.4f} 秒")