faiss-gpu-cu12

import faiss import numpy as np # 1. Define dataset parameters d = 128 # dimension nb = 100000 # database size nq = 10 # number of queries # 2. Generate random data np.random.seed(1234) xb = np.random.random((nb, d)).astype('float32') xq = np.random.random((nq, d)).astype('float32') # Ensure the data is C-contiguous as Faiss often expects it xb = np.ascontiguousarray(xb) xq = np.ascontiguousarray(xq) # 3. Build a CPU index (e.g., L2 distance) index_cpu = faiss.IndexFlatL2(d) print(f"Is CPU index trained? {index_cpu.is_trained}") # 4. Add vectors to the CPU index index_cpu.add(xb) print(f"Number of vectors in CPU index: {index_cpu.ntotal}") # 5. Attempt to move the index to GPU try: # Faiss GPU indices require StandardGpuResources res = faiss.StandardGpuResources() # 0 for the first GPU; change if you have multiple and want a different one index_gpu = faiss.index_cpu_to_gpu(res, 0, index_cpu) print(f"\nIndex successfully moved to GPU. Number of vectors: {index_gpu.ntotal}") # 6. Perform search on GPU k = 4 # We want to find 4 nearest neighbors D, I = index_gpu.search(xq, k) # D for distances, I for indices print("\nGPU Search - Distances (D):") print(D) print("\nGPU Search - Indices (I):") print(I) except Exception as e: print(f"\nCould not run GPU operations. This might happen if no compatible GPU or drivers are found, or if it's a CPU-only environment. Error: {e}") print("Falling back to CPU search for demonstration:") k = 4 D_cpu, I_cpu = index_cpu.search(xq, k) print("\nCPU Search - Distances (D):") print(D_cpu) print( "\nCPU Search - Indices (I):") print(I_cpu)