如何测试faiss是否为gpu版本

最新推荐文章于 2024-05-20 11:52:24 发布

遇到好事了

最新推荐文章于 2024-05-20 11:52:24 发布

阅读量606

点赞数 3

文章标签： faiss

本文链接：https://blog.csdn.net/qq_42549774/article/details/138192133

版权

本文通过实测比较了Faiss在支持GPU和不支持GPU情况下的搜索时间，结果显示GPU版本搜索时间显著降低，验证了Faiss在GPU上的有效性能提升。

摘要由CSDN通过智能技术生成

这个脚本将测量使用GPU进行搜索的时间。如果搜索时间显著减少，那么GPU版本的Faiss应该正常工作。
本人实测，gpu版本的输出应该如下：

Faiss支持GPU
搜索时间: 0.0423 秒

cpu的输出为：

Faiss不支持GPU
搜索时间: 66.4641 秒

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} 秒")