faiss_cluster_cpu.py

s = time.time()

k = 10
n_init = 10
max_iter = 300
kmeans = faiss.Kmeans(d=x_train.shape[1], k=k, niter=max_iter, nredo=n_init)
kmeans.train(x_train.astype(np.float32))

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


s = time.time()

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

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