demo_independent_ivf_dimension.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "641bbab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import faiss\n",
    "\n",
    "from faiss.contrib.datasets import DatasetDeep1B\n",
    "from faiss.contrib.ivf_tools import add_preassigned, search_preassigned"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a93f105",
   "metadata": {},
   "source": [
    "This script demonstrates how to use a high-dimensional coarse quantizer with a low-dimensional fine quantizer. \n",
    "This is not possible out-of-the-box because the IVFPQ implementation assumes the IVF quantizer and the PQ run in the same dimension. \n",
    "To combine the quantizers in different dimensionalities, the approach is to use `search_preassigned` and `add_preassigned` to perform the coarse quantization and add / search separately.  \n",
    "\n",
    "In this example, the OPQ pre-transformation (an orthonormal transformation of the data) reduces the dimension of the input from 96 to 32 dimensions, so the coarse quantizer may not be as selective as it could.\n",
    "By doing the coarse quantization and the search separately, the accuracy improves for some (but not all) settings of `nprobe`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "67691821",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset in dimension 96, with metric L2, size: Q 10000 B 1000000 T 358480000\n"
     ]
    }
   ],
   "source": [
    "ds = DatasetDeep1B(nb=10**6)\n",
    "gt = ds.get_groundtruth()\n",
    "\n",
    "print(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8cf7944",
   "metadata": {},
   "source": [
    "# Baseline "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3ea1e7e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.index_factory(ds.d, \"OPQ4_32,IVF4096,PQ4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "763de8f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.train(ds.get_train(maxtrain=10**6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c193c307",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.add(ds.get_database())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9e9ceed6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nprobe=  1  recalls={1: 0.0793, 10: 0.2318, 100: 0.3194}\n",
      "nprobe=  4  recalls={1: 0.0983, 10: 0.3433, 100: 0.58}\n",
      "nprobe= 16  recalls={1: 0.1026, 10: 0.3901, 100: 0.7409}\n"
     ]
    }
   ],
   "source": [
    "for nprobe in 1, 4, 16: \n",
    "    faiss.extract_index_ivf(index).nprobe = nprobe \n",
    "    D, I = index.search(ds.get_queries(), 100)\n",
    "    recalls = {\n",
    "        rank: (I[:, :rank] == gt[:, :1]).sum() / len(gt)\n",
    "        for rank in [1, 10, 100]        \n",
    "    }\n",
    "    print(f\"{nprobe=:3d}  {recalls=:}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22892c83",
   "metadata": {},
   "source": [
    "## Decompose computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9f966129",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.reset()\n",
    "index_ivf = faiss.extract_index_ivf(index)\n",
    "preproc = index.chain.at(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5f1531dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb = preproc.apply(ds.get_database())\n",
    "# coarse quantization\n",
    "Dc, Ic = index_ivf.quantizer.search(xb, 1)\n",
    "# add operation \n",
    "add_preassigned(index_ivf, xb, Ic.ravel())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "23a022d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nprobe=  1  recalls={1: 0.0793, 10: 0.2318, 100: 0.3194}\n",
      "nprobe=  4  recalls={1: 0.0983, 10: 0.3433, 100: 0.58}\n",
      "nprobe= 16  recalls={1: 0.1026, 10: 0.3901, 100: 0.7409}\n"
     ]
    }
   ],
   "source": [
    "for nprobe in 1, 4, 16: \n",
    "    index_ivf.nprobe = nprobe \n",
    "    xq = preproc.apply(ds.get_queries())\n",
    "        \n",
    "    # coarse quantization \n",
    "    Dc, Ic = index_ivf.quantizer.search(xq, nprobe)\n",
    "    \n",
    "    # actual search \n",
    "    D, I = search_preassigned(index_ivf, xq, 100, Ic, Dc)\n",
    "    recalls = {\n",
    "        rank: (I[:, :rank] == gt[:, :1]).sum() / len(gt)\n",
    "        for rank in [1, 10, 100]        \n",
    "    }\n",
    "    print(f\"{nprobe=:3d}  {recalls=:}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85655978",
   "metadata": {},
   "source": [
    "So here we perform exactly the same operation, but with an explicit separation of the coarse quantizer."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72af31c4",
   "metadata": {},
   "source": [
    "# Test with full-dimensional quantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "4657cb1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.index_factory(ds.d, \"OPQ4_32,IVF4096,PQ4\")\n",
    "qindex = faiss.index_factory(ds.d, \"IVF4096,Flat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "0b9f1a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the full-dimensional \"big\" coarse quantizer\n",
    "qindex.train(ds.get_train(maxtrain=10**6))\n",
    "quantizer = qindex.quantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "939d3bfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_ivf = faiss.downcast_index(faiss.extract_index_ivf(index))\n",
    "preproc = index.chain.at(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "79befc7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train the OPQ \n",
    "xt = ds.get_train(maxtrain=10**6)\n",
    "preproc.train(xt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "1a72fdf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "xt2 = preproc.apply(xt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "31989a66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare the \"small\" coarse quantizer \n",
    "centroids = quantizer.reconstruct_n(0, quantizer.ntotal)\n",
    "centroids2 = preproc.apply(xt)\n",
    "index_ivf.quantizer.add(centroids2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "82dd943a",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_ivf.train(xt2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "530de9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.is_trained = True "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "ffd995d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "xb = ds.get_database()\n",
    "xb2 = preproc.apply(xb)\n",
    "# coarse quantization with the big quantizer\n",
    "Dc, Ic = quantizer.search(xb, 1)\n",
    "# add operation \n",
    "add_preassigned(index_ivf, xb2, Ic.ravel())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "b12792f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# disable precomputed tables, because the Dc is out of sync with the \n",
    "# small coarse quantizer\n",
    "index_ivf.use_precomputed_table = -1\n",
    "index_ivf.precompute_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "4b044a90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nprobe=  1  recalls={1: 0.0831, 10: 0.2694, 100: 0.387}\n",
      "nprobe=  4  recalls={1: 0.0873, 10: 0.3341, 100: 0.634}\n",
      "nprobe= 16  recalls={1: 0.0851, 10: 0.3415, 100: 0.7189}\n"
     ]
    }
   ],
   "source": [
    "for nprobe in 1, 4, 16: \n",
    "    index_ivf.nprobe = nprobe \n",
    "    \n",
    "    xq = ds.get_queries()\n",
    "    xq2 = preproc.apply(xq)\n",
    "        \n",
    "    # coarse quantization \n",
    "    Dc, Ic = quantizer.search(xq, nprobe)\n",
    "    nq, d = xq2.shape\n",
    "    \n",
    "    # actual search \n",
    "    D, I = search_preassigned(index_ivf, xq2, 100, Ic, None)\n",
    "    recalls = {\n",
    "        rank: (I[:, :rank] == gt[:, :1]).sum() / len(gt)\n",
    "        for rank in [1, 10, 100]        \n",
    "    }\n",
    "    print(f\"{nprobe=:3d}  {recalls=:}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7c58343",
   "metadata": {},
   "source": [
    "Better than the baseline for low nprobe values, but getting worse for higher values. "
   ]
  }
 ],
 "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}