Skip to content

Running on GPUs

Lucas Hosseini edited this page Feb 21, 2018 · 8 revisions

Faiss can leverage your nvidia GPUs.

In Python

print("number of GPUs:", faiss.get_num_gpus())

res = faiss.StandardGpuResources()  # use a single GPU

index_flat = faiss.IndexFlatL2(d)   # build a flat (CPU) index

# make it a GPU index
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)

gpu_index_flat.add(xb)         # add vectors to the index
print(gpu_index_flat.ntotal)

k = 4                          # we want to see 4 nearest neighbors
D, I = gpu_index_flat.search(xq, k)  # actual search
print(I[:5])                   # neighbors of the 5 first queries
print(I[-5:])                  # neighbors of the 5 last queries

In C++

    faiss::gpu::StandardGpuResources res; // use a single GPU

    faiss::gpu::GpuIndexFlatL2 index_flat(&res, d); // build a flat GPU index

    printf("is_trained = %s\n", index_flat.is_trained ? "true" : "false");
    index_flat.add(nb, xb);  // add vectors to the index
    printf("ntotal = %ld\n", index_flat.ntotal);

    int k = 4;
    {       // search xq
        long *I = new long[k * nq];
        float *D = new float[k * nq];

        index_flat.search(nq, xq, k, D, I);

        // print results
        printf("I (5 first results)=\n");
        for(int i = 0; i < 5; i++) {
            for(int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        printf("I (5 last results)=\n");
        for(int i = nq - 5; i < nq; i++) {
            for(int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("\n");
        }

        delete [] I;
        delete [] D;
    }
Clone this wiki locally