Moving a basic CPU index to GPU (Google Colab T4) for faster search?

[neetscience]

Hi everyone,

I am new to this field and wanted to know if there is a straightforward way to accelerate a standard CPU index (like IndexFlatL2) using a GPU.

I built a basic index that works perfectly on my CPU, but since I have access to a Google Colab T4, I'd love to take advantage of it for faster queries. Is it a complex process to transition the memory and search operations over? And how do I actually format the search to see my nearest neighbors once it is on the GPU?

[BryanBradfo]

Welcome @neetscience to vector search! The great news is that transitioning to the GPU in Faiss is incredibly simple. You don't need to rewrite your indexing logic—Faiss has a built-in helper function called index_cpu_to_gpu that handles copying the vectors from your system memory over to the GPU's memory.

Here is a complete, concise example that runs perfectly in a Google Colab T4 instance. Note that we are installing faiss-gpu-cu12 since Colab uses CUDA 12!

# 1. Install the modern CUDA 12 package for Google Colab
!pip install faiss-gpu-cu12

import faiss
import numpy as np

# 2. Generate dummy database and query vectors
dimension = 128
database_size = 10000
xb = np.random.random((database_size, dimension)).astype('float32') 
# 5 new queries
xq = np.random.random((5, dimension)).astype('float32')

# 3. Build your standard CPU index
cpu_index = faiss.IndexFlatL2(dimension)
cpu_index.add(xb)

# 4. The Magic Step: Transfer it to the T4 GPU
res = faiss.StandardGpuResources() # Initialize GPU resources
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) # 0: your GPU ID

# 5. Perform the search directly on the GPU
k = 4 # Find the top 4 closest matches
distances, indices = gpu_index.search(xq, k)

print("Which vectors matched? (Row IDs in our database):")
print(indices)
print("\nHow close were they? (L2 Distances):")
print(distances)

Understanding the output:

• indices: Returns a 5x4 matrix. These are the exact row IDs of the matching vectors in your original xb database.
• distances: Returns the mathematical L2 (Euclidean) distance between your query and the matches. Smaller numbers mean the vectors are closer/more similar!

By keeping both the index and the search operations on the GPU, you'll see a massive speedup when your dataset grows from 10,000 vectors to 10 million!

[neetscience]

Thank you so much, @BryanBradfo! This is exactly what I was looking for.

The code ran perfectly, and it is amazing how seamless index_cpu_to_gpu makes the whole transition. Seeing the distances and indices output makes total sense now.