faiss-gpu

import faiss import numpy as np # Check for GPU availability if not faiss.get_num_gpus(): print("Warning: No GPUs detected. Faiss GPU functionality will not be available.") print("Ensure NVIDIA drivers and CUDA toolkit are correctly installed.") # In a real application, you might exit or switch to faiss-cpu here. # For this quickstart, we will proceed assuming GPU context is desired if available. # Define parameters d = 128 # vector dimension nb = 100000 # database size nq = 1000 # number of query vectors k = 4 # number of nearest neighbors to search for # Generate random data np.random.seed(1234) xb = np.random.random((nb, d)).astype('float32') xb[:, 0] += np.arange(nb) / 1000. xq = np.random.random((nq, d)).astype('float32') xq[:, 0] += np.arange(nq) / 1000. print(f"Database vectors shape: {xb.shape}") print(f"Query vectors shape: {xq.shape}") # Initialize GPU resources # A single StandardGpuResources object handles memory management on a specific GPU. res = faiss.StandardGpuResources() # Create a GPU index (e.g., IndexFlatL2 for brute-force L2 distance) # 'res' links the index to the GPU resources, 'd' is the vector dimension. index_flat = faiss.GpuIndexFlatL2(res, d) # Add vectors to the index print("Adding vectors to GPU index...") index_flat.add(xb) print(f"Number of vectors in the index: {index_flat.ntotal}") # Search for nearest neighbors print(f"Searching for {k} nearest neighbors for {nq} queries...") D, I = index_flat.search(xq, k) # D: distances, I: indices print("\nFirst 5 query results (indices of nearest neighbors):") print(I[:5]) print("\nFirst 5 query results (distances):") print(D[:5])