tutorial/4_comparison_to_faiss.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 4: Comparison to faiss\n",
    "This chapter contains the followings:\n",
    "\n",
    "1. Setup the experiment using SIFT1M\n",
    "1. Small-scale comparison: N=10^5, K=10^3 (k-means with faiss-CPU and k-means with sklearn)\n",
    "1. Large-scale comparison: N=10^6, K=10^4 (PQk-means, k-means with faiss-CPU, and k-measn with falss-GPU)\n",
    "\n",
    "Requisites:\n",
    "- numpy\n",
    "- pqkmeans\n",
    "- sklearn\n",
    "- [faiss](https://github.com/facebookresearch/faiss) (you can install it via conda)\n",
    "  1. CPU version: `conda install faiss-cpu -c pytorch`\n",
    "  1. GPU version (with two NVIDIA GTX1080s): `conda install faiss-gpu -c pytorch`\n",
    "  \n",
    "  \n",
    "Our final suggestions are as follows:\n",
    "- If you have GPU(s) and your GPU memory is large enough (all data can be loaded on the GPU memory at once), faiss-GPU is the fastest option.\n",
    "- Otherwise, \n",
    "  - If your problem is small enough (all vectors can be easily fit into the RAM), faiss-CPU would be the best option.\n",
    "  - If the problem is large, e.g., (1) faiss-CPU seems slow, or (2) the vectors cannot be loaded on the memory at once, then PQk-means is the best option.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup the experiment using SIFT1M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "import pqkmeans\n",
    "from sklearn.cluster import KMeans\n",
    "import faiss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this chapter, we compare our PQk-means to k-means in the [faiss](https://github.com/facebookresearch/faiss) library. Faiss provides one of the most efficient implementations of nearest neighbor algorithms for both CPU(s) and GPU(s). It also provides an implementation of vanilla k-means, which we will compare to. The core part of faiss is implemented by C++, and the python binding is available.\n",
    "\n",
    "We compare PQk-means to both CPU- and GPU-version. Our configurations are:\n",
    "- faiss-CPU: This was built with Intel MKL, which provides the fastest backend BLAS implementation. The algorithms in the library are automatically parallelized. All evaluations are conducted on a server with 3.6 GHz Intel Xeon CPU (6 cores, 12 threads)\n",
    "- faiss-GPU: The library was built with CUDA 8.0. Two middle-level GPUs, NVIDIA GTX 1080s, are used for the evaluation. The algorithms can be run over multi GPUs. \n",
    "\n",
    "For the comparison, we leverage the SIFT1M dataset.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xt, X = pqkmeans.evaluation.get_sift1m_dataset()  #  Xt: the training data.  X: the testing data to be clusterd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, you can download the data by a helper script. This would take several minutes, and consume 168 MB of the disk space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Xt.shape:(100000, 128)\n",
      "X.shape:(1000000, 128)\n"
     ]
    }
   ],
   "source": [
    "Xt = Xt.astype(numpy.float32)\n",
    "X = X.astype(numpy.float32)\n",
    "D = X.shape[1]\n",
    "print(\"Xt.shape:{}\\nX.shape:{}\".format(Xt.shape, X.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because faiss takes 32-bit float vectors as inputs, the data is converted to float32."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Small-scale comparison: N=10^5, K=10^3 (k-means with faiss-CPU v.s. k-means with sklearn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let us compare the k-means implementation of faiss and sklearn using 100K vectors from SIFT1M. Then we show that faiss is much faster than sklearn with almost the same error.\n",
    "\n",
    "Note that it is hard to run k-means-sklearn with a large K because it is too slow (that is the reason for this small-scale experiment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "K_small = 1000\n",
    "N_small = 100000\n",
    "\n",
    "# Setup clustering instances. We stop each algorithm with 10 iterations\n",
    "kmeans_faiss_cpu_small = faiss.Kmeans(d=D, k=K_small, niter=10)\n",
    "kmeans_sklearn_small = KMeans(n_clusters=K_small, n_jobs=-1, max_iter=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run each algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "faiss-cpu:\n",
      "CPU times: user 18.1 s, sys: 1.82 s, total: 20 s\n",
      "Wall time: 1.8 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"faiss-cpu:\")\n",
    "kmeans_faiss_cpu_small.train(X[:N_small])\n",
    "_, ids_faiss_cpu_small = kmeans_faiss_cpu_small.index.search(X[:N_small], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sklearn:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n",
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/sklearn/metrics/pairwise.py:257: RuntimeWarning: invalid value encountered in sqrt\n",
      "  return distances if squared else np.sqrt(distances, out=distances)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.86 s, sys: 380 ms, total: 3.24 s\n",
      "Wall time: 2min 53s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"sklearn:\")\n",
    "ids_sklearn_small = kmeans_sklearn_small.fit_predict(X[:N_small])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k-means, faiss-cpu, error: 216.8936472830963\n",
      "k-means, sklearn, error: 216.3501130987549\n"
     ]
    }
   ],
   "source": [
    "_, faiss_cpu_small_error, _ = pqkmeans.evaluation.calc_error(ids_faiss_cpu_small.reshape(-1), X[:N_small], K_small)\n",
    "_, sklearn_small_error, _ = pqkmeans.evaluation.calc_error(ids_sklearn_small, X[:N_small], K_small)\n",
    "\n",
    "print(\"k-means, faiss-cpu, error: {}\".format(faiss_cpu_small_error))\n",
    "print(\"k-means, sklearn, error: {}\".format(sklearn_small_error))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observed that\n",
    "- k-means with faiss-CPU (2 sec) is surprisingly faster than k-means with sklearn (3 min) with almost the same error. This speedup would be due to the highly optimized implementation of the nearest neighbor search in faiss with Intel MKL BLAS. This suggests that faiss-CPU is a better option for the exact k-means in a usual computer.\n",
    "\n",
    "Because faiss-CPU is faster thant sklearn, sklearn is not compared in the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Large-scale comparison: N=10^6, K=10^4 (PQk-means, k-means with faiss-CPU, and k-measn with falss-GPU)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup GPUs for faiss-gpu\n",
    "# In my environment, the first GPU (id=0) is for rendering, and the second (id=1) and the third (id=2) GPUs are GPGPU (GTX1080).\n",
    "# We activate only the second and the third GPU\n",
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"  # make sure the order is identical to the result of nvidia-smi\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1,2\"  # Please change here for your environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let us compare PQk-meas with faiss-CPU and faiss-GPU using the whole dataset (N=10^6, K=10^4). Note that this is 100x larter setting compared to Sec 2 (NK=10^8 vs NK=10^10). \n",
    "\n",
    "First, as pre-processing for PQk-means, let's train a PQencoder and encode all data. It will take around 10 sec."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matsui/usr/src/anaconda/anaconda3/envs/pqkmeans_faiss/lib/python3.6/site-packages/scipy/cluster/vq.py:523: UserWarning: One of the clusters is empty. Re-run kmeans with a different initialization.\n",
      "  warnings.warn(\"One of the clusters is empty. \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 57.3 s, sys: 7.74 s, total: 1min 5s\n",
      "Wall time: 11 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Train the encoder\n",
    "encoder = pqkmeans.encoder.PQEncoder(num_subdim=4, Ks=256)\n",
    "encoder.fit(Xt)\n",
    "\n",
    "# Encode the vectors to PQ-codes\n",
    "X_code = encoder.transform(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that X_code is 128x more memory efficient than X:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X.shape: (1000000, 128), X.dtype: float32, X.nbytes: 512.0 MB\n",
      "X_code.shape: (1000000, 4), X_code.dtype: uint8, X_code.nbytes: 4.0 MB\n"
     ]
    }
   ],
   "source": [
    "print(\"X.shape: {}, X.dtype: {}, X.nbytes: {} MB\".format(X.shape, X.dtype, X.nbytes / 10**6))\n",
    "print(\"X_code.shape: {}, X_code.dtype: {}, X_code.nbytes: {} MB\".format(X_code.shape, X_code.dtype, X_code.nbytes / 10**6))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then each algorithms are instantiated as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 10000  # Set larger K\n",
    "\n",
    "# Setup k-means instances. The number of iteration is set as 20 for all methods\n",
    "\n",
    "# PQ-kmeans\n",
    "kmeans_pqkmeans = pqkmeans.clustering.PQKMeans(encoder=encoder, k=K, iteration=20)\n",
    "\n",
    "# Faiss-cpu\n",
    "kmeans_faiss_cpu = faiss.Kmeans(d=D, k=K, niter=20)\n",
    "kmeans_faiss_cpu.cp.max_points_per_centroid = 1000000   #  otherwise the kmeans implementation sub-samples the training set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because some configurations are required for GPU, we wrap up the gpu clustering as one function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_faiss_gpu(X, K, ngpu):\n",
    "    # This code is based on https://github.com/facebookresearch/faiss/blob/master/benchs/kmeans_mnist.py\n",
    "    \n",
    "    D = X.shape[1]\n",
    "    clus = faiss.Clustering(D, K)\n",
    "    \n",
    "    # otherwise the kmeans implementation sub-samples the training set\n",
    "    clus.max_points_per_centroid = 10000000\n",
    "    \n",
    "    clus.niter = 20\n",
    "    \n",
    "    res = [faiss.StandardGpuResources() for i in range(ngpu)]\n",
    "\n",
    "    flat_config = []\n",
    "    for i in range(ngpu):\n",
    "        cfg = faiss.GpuIndexFlatConfig()\n",
    "        cfg.useFloat16 = False\n",
    "        cfg.device = i\n",
    "        flat_config.append(cfg)\n",
    "\n",
    "    if ngpu == 1:\n",
    "        index = faiss.GpuIndexFlatL2(res[0], D, flat_config[0])\n",
    "    else:\n",
    "        indexes = [faiss.GpuIndexFlatL2(res[i], D, flat_config[i])\n",
    "                   for i in range(ngpu)]\n",
    "        index = faiss.IndexProxy()\n",
    "        for sub_index in indexes:\n",
    "            index.addIndex(sub_index)\n",
    "            \n",
    "    \n",
    "    # Run clustering\n",
    "    clus.train(X, index)\n",
    "    \n",
    "    # Return the assignment\n",
    "    _, ids = index.search(X, 1)\n",
    "    return ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run each method and see the computational cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PQk-means:\n",
      "CPU times: user 29min 5s, sys: 3.45 s, total: 29min 8s\n",
      "Wall time: 2min 34s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"PQk-means:\")\n",
    "ids_pqkmeans = kmeans_pqkmeans.fit_predict(X_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "faiss-cpu:\n",
      "CPU times: user 48min 9s, sys: 2min 43s, total: 50min 53s\n",
      "Wall time: 4min 21s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"faiss-cpu:\")\n",
    "kmeans_faiss_cpu.train(X)\n",
    "_, ids_faiss_cpu = kmeans_faiss_cpu.index.search(X, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "faiss with GPU:\n",
      "CPU times: user 1min, sys: 5.44 s, total: 1min 5s\n",
      "Wall time: 10.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print(\"faiss with GPU:\")\n",
    "ids_faiss_gpu = run_faiss_gpu(X, K, ngpu=2) # Please adjust ngpu for your environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PQk-means, error: 218.5069038487091\n",
      "k-means, faiss-cpu, error: 197.61731756228255\n",
      "k-means, faiss-gpu, error: 197.61766727244188\n"
     ]
    }
   ],
   "source": [
    "_, pqkmeans_error, _ = pqkmeans.evaluation.calc_error(ids_pqkmeans, X, K)\n",
    "_, faiss_cpu_error, _ = pqkmeans.evaluation.calc_error(ids_faiss_cpu.reshape(-1), X, K)\n",
    "_, faiss_gpu_error, _ = pqkmeans.evaluation.calc_error(ids_faiss_gpu.reshape(-1), X, K)\n",
    "\n",
    "print(\"PQk-means, error: {}\".format(pqkmeans_error))\n",
    "print(\"k-means, faiss-cpu, error: {}\".format(faiss_cpu_error))\n",
    "print(\"k-means, faiss-gpu, error: {}\".format(faiss_gpu_error))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our observations are:\n",
    "- PQk-means (around 2 min) is 2x faster than k-means with faiss-CPU (around 4 min). The cost of learining/encoding is marginal (10 sec). \n",
    "- PQk-means is memory efficient (128x in this case). More data can be easily processed even if the data itself cannot be loaded on the RAM at once (see [tutorial3](./3_billion_scale_clustering.ipynb)). Note that faiss provides memory-efficient search algorithms including IVFPQ, but the clustering itself is vanilla k-means (all original vectors need to be loaded on the memory).\n",
    "- Because PQk-means is an approximation of k-means, the accuracy of clustering is lower than k-means with CPU/GPU faiss.\n",
    "- k-means with faiss-GPU (10 sec) is suprisingly faster than both PQk-means and faiss-CPU, with the same error as faiss-CPU. We conlude that, if you have several GPUs, faiss-GPU is the fastest option for the exact k-means (see [benchmark](https://github.com/facebookresearch/faiss/wiki/Low-level-benchmarks) for more results). Note that PQk-means with GPUs could be faster, but has not been implemented yet.\n",
    "\n",
    "\n",
    "Our final suggestions are as follows:\n",
    "- If you have GPU(s) and your GPU memory is large enough (all data can be loaded on the GPU memory at once), faiss-GPU is the fastest option.\n",
    "- Otherwise, \n",
    "  - If your problem is small enough (all vectors can be easily fit into the RAM), faiss-CPU would be the best option.\n",
    "  - If the problem is large, e.g., (1) faiss-CPU seems slow, or (2) the vectors cannot be loaded on the memory at once, then PQk-means is the best option."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}