Skip to content

Instantly share code, notes, and snippets.

@GeorgeSeif

GeorgeSeif/faiss_cluster_gpu.py

Created January 30, 2021 18:15
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save GeorgeSeif/d0917dee46724e950e00db9ebf6f8b6f to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save GeorgeSeif/d0917dee46724e950e00db9ebf6f8b6f to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. GeorgeSeif created this gist Jan 30, 2021.
    16 changes: 16 additions & 0 deletions faiss_cluster_gpu.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,16 @@
    k = 10
    n_init = 10
    max_iter = 300
    kmeans = faiss.Kmeans(d=data.shape[1], k=k, niter=max_iter, nredo=n_init, gpu=True)
    kmeans.train(data.astype(np.float32))

    e = time.time()
    print("Training time = {}".format(e - s))


    s = time.time()

    kmeans.index.search(data.astype(np.float32), 1)[1]

    e = time.time()
    print("Prediction time = {}".format((e - s) / data_size))