Faiss GPU

Warnings

• breaking The `faiss-gpu` PyPI package (v1.7.2, released early 2022) is significantly outdated compared to the upstream Faiss library (latest official v1.9.1; `faiss-wheels` targets v1.13.x). This means many new features, bug fixes, and performance optimizations from more recent Faiss versions will not be available, and API signatures may differ from current documentation. For newer Faiss versions with GPU support, users often need to compile from source or seek alternative, potentially unofficial, wheel distributions.
  Fix: For the latest Faiss features, consider compiling Faiss from source (from `facebookresearch/faiss`) with GPU support. If a newer `faiss-gpu` PyPI package eventually becomes available, upgrade to it. Alternatively, `faiss-cpu` offers a more up-to-date PyPI presence for CPU-only workflows.
• gotcha Using `faiss-gpu` requires a properly installed NVIDIA CUDA toolkit, compatible NVIDIA drivers, and potentially cuBLAS/cuDNN libraries. The Faiss build is linked against specific CUDA versions. Missing or incompatible GPU dependencies will lead to runtime errors (e.g., `FaissError` during index creation or `ModuleNotFoundError` if GPU-specific components cannot load).
  Fix: Refer to the Faiss installation guides and your NVIDIA CUDA documentation for your specific OS and hardware. Ensure `nvcc --version` shows a compatible CUDA toolkit. On Linux, ensure `LD_LIBRARY_PATH` (or equivalent) correctly points to your CUDA library directories.
• gotcha While the PyPI package name is `faiss-gpu`, the Python module is imported as `faiss`. Attempting `import faiss_gpu` will result in an `ImportError`. GPU-specific functionality is then accessed through methods and classes within the `faiss` module (e.g., `faiss.GpuIndexFlatL2`).
  Fix: Always use `import faiss` in your Python code.
• gotcha Compatibility with NumPy 2.0 (released mid-2024) is a known issue for many C-extension libraries, including Faiss. Older `faiss-gpu` versions (like 1.7.2) are unlikely to be compatible without recompilation or patches. The `faiss-wheels` project, which typically builds these wheels, explicitly pins NumPy versions to `<2.0` in its build configurations.
  Fix: To avoid potential runtime errors, ensure your environment uses a NumPy version less than 2.0 (e.g., `pip install 'numpy<2.0'`).

Install

• pip install faiss-gpu Install `faiss-gpu`

Imports

• faiss
  wrong:

  import faiss_gpu

  correct:

  import faiss

  The Python package name is `faiss-gpu`, but the module is imported as `faiss`. GPU-specific classes like `GpuIndexFlatL2` are then accessed through the `faiss` module.

Quickstart

This quickstart demonstrates how to initialize a GPU-accelerated Faiss index, add vectors to it, and perform a similarity search. It uses `faiss.StandardGpuResources` to manage GPU memory and `faiss.GpuIndexFlatL2` for a basic brute-force L2 distance search. Ensure your system has a compatible NVIDIA GPU and CUDA toolkit installed.

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])