如何测试faiss是否为gpu版本

遇到好事了2024-04-29 9:43

这个脚本将测量使用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} 秒")