faiss_gpu_index_transfer_knn_search_quickstart.py

1import numpy as np
2import faiss
3
4# Dimension of the vectors
5d = 64                           
6# Number of database vectors
7nb = 100000                      
8# Number of query vectors
9nq = 10000                       
10np.random.seed(1234)             
11
12# Generate synthetic data
13xb = np.random.random((nb, d)).astype('float32')
14xb[:, 0] += np.arange(nb) / 1000.
15xq = np.random.random((nq, d)).astype('float32')
16xq[:, 0] += np.arange(nq) / 1000.
17
18# Build a flat (L2) index on the CPU
19index_flat = faiss.IndexFlatL2(d)
20
21# Initialize the GPU resources
22res = faiss.StandardGpuResources()
23
24# Transfer the index to the GPU
25gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
26
27# Add vectors to the index
28gpu_index_flat.add(xb)         
29print(f"Index size: {gpu_index_flat.ntotal}")
30
31# Search for the 4 nearest neighbors
32k = 4                          
33D, I = gpu_index_flat.search(xq, k) 
34
35# Print results
36print("First 5 search results (indices):")
37print(I[:5])
38print("First 5 search results (distances):")
39print(D[:5])