Faiss-GPU library import
To use Faiss-GPU, you need to import torch before faiss. Otherwise, you will get an error like this when you are trying to use faiss gpu_index = faiss.index_cpu_to_gpu(res, 0, index) # make it a GPU index:
1
RuntimeError: Error in faiss::gpu::GpuIndex::GpuIndex(std::shared_ptr, int, faiss::MetricType, float, faiss::gpu::GpuIndexConfig) at /root/miniconda3/conda-bld/faiss-pkg_1669821803039/work/faiss/gpu/GpuIndex.cu:55: Error: 'config_.device < getNumDevices()' failed: Invalid GPU device 0
and also when trying to use a = torch.tensor([[1.,2,3],[4,5,6]], device='cuda')
1
RuntimeError: No CUDA GPUs are available
Solution: import torch before faiss:
1
2
3
4
import torch
import faiss
import faiss.contrib.torch_utils
#import torch
Example code for Faiss-GPU
IndexFlatL2
Note that the distance here is the squared Euclidean (L2) distance, avoiding the square root. If we want the exact distance between two features, we need to apply square root.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
index = faiss.IndexFlatL2(3) # build the index
res = faiss.StandardGpuResources() # use a single GPU
gpu_index = faiss.index_cpu_to_gpu(res, 0, index) # make it a GPU index
a = torch.tensor([[1.,2,3],[4,5,6],[4,5,100], [4,123,6], [434,5,6]], device='cuda')
b = torch.tensor([[1.,2.,3.]], device='cuda')
# Normalize before adding to index
doc = f.normalize(a, p=2, dim=1)
q = f.normalize(b, p=2, dim=1)
gpu_index.add(doc)
D,I = gpu_index.search(q, topk)
# D gives the L2 distacne
# I gives index, the closest(smallest D)/most similar index in the docs first
L2squared_pytorch = torch.nn.MSELoss(reduction='none')
L2_squared = L2squared_pytorch(a,b)
L2_squared = L2_squared.sum(dim=1)
# at this point L2_squared should be consistent as to D
# if you want to get actual distance:
L2_squared = L2_squared ** 0.5
IndexFlatIP
Inner product, if normalized then we would have a result equivalent to cosine similarity. Note here normalised Inner product measures similarity (range +1 → -1), larger the value, more similar which is the opposite to the distance metrics in L2.
1
2
3
4
5
6
7
8
9
10
11
12
13
index_IP = faiss.IndexFlatIP(3)
res = faiss.StandardGpuResources() # use a single GPU
gpu_index_IP = faiss.index_cpu_to_gpu(res, 0, index_IP) # make it a GPU index
gpu_index_IP.add(doc)
Sim,I = gpu_index_IP.search(q, topk)
# Sim gives the similarity
# I gives index, the closest/most similar(largest Sim) index in the docs first
# Given that the embeddings are normalised, Sim should be equivalent to cosine sim
cos = torch.nn.functional.cosine_similarity(a,b)
# cos should be approximately equal to Sim,
# Due to different precision, they are not exactly equal in floats
Reference
- https://github.com/facebookresearch/faiss/wiki/MetricType-and-distances
- https://github.com/facebookresearch/faiss/wiki/Faiss-indexes
- https://pytorch.org/docs/master/generated/torch.nn.MSELoss.html#mseloss
- https://www.facebook.com/groups/faissusers/posts/1025663204524632/
- https://github.com/facebookresearch/faiss/wiki/Faiss-on-the-GPU
- https://github.com/facebookresearch/faiss/wiki/Faiss-on-the-GPU
