Choose FAISS gpu usage #292
Description
Hi,
Thanks a lot for the library !
I had an issue using AccuracyCalculator because FAISS was using other gpus than the one I was working on.
Would it be possible to change the function get_knn so that it would take an optional argument, that if passed the FAISS knn search would be done on a single gpu (the one where the embeddings are stored).
def get_knn(
reference_embeddings, test_embeddings, k, embeddings_come_from_same_source=False, use_all_gpu=True
):
device = reference_embeddings.device
reference_embeddings = c_f.to_numpy(reference_embeddings).astype(np.float32)
test_embeddings = c_f.to_numpy(test_embeddings).astype(np.float32)
d = reference_embeddings.shape[1]
logging.info("running k-nn with k=%d" % k)
logging.info("embedding dimensionality is %d" % d)
index = faiss.IndexFlatL2(d)
if faiss.get_num_gpus() > 0:
if use_all_gpu:
index = faiss.index_cpu_to_all_gpus(index)
else:
res = faiss.StandardGpuResources()
index = faiss.index_cpu_to_gpu(res, device.index, index)
index.add(reference_embeddings)
distances, indices = index.search(test_embeddings, k + 1)
distances = c_f.to_device(torch.from_numpy(distances), device=device)
indices = c_f.to_device(torch.from_numpy(indices), device=device)
if embeddings_come_from_same_source:
return indices[:, 1:], distances[:, 1:]
return indices[:, :k], distances[:, :k]
The following changes would have to be also done for the run_kmeans function.
And then you would be able to pass the same argument to AccuracyCalculator.
Thanks