assign_on_gpu.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import faiss\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare an index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 128\n",
    "nt = 10000\n",
    "nb = 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs = np.random.RandomState(123)\n",
    "\n",
    "xbt = rs.rand(nt + nb, d).astype('float32')\n",
    "xb = xbt[:nb]\n",
    "xt = xbt[nb:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = faiss.index_factory(d, 'PCA64,IVF1024,PQ16')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "index.train(xt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Populate index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get original quantizer \n",
    "index_ivf = faiss.extract_index_ivf(index)\n",
    "quantizer = index_ivf.quantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# move quantizer to GPU\n",
    "quantizer_gpu = faiss.index_cpu_to_all_gpus(quantizer, ngpu=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from this point on quantizer_gpu should remain in scope! Python does not track this reference.\n",
    "index_ivf.quantizer = quantizer_gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quantization is performed on GPU\n",
    "# any number of add() and add_with_ids() calls can go here\n",
    "index.add(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_ivf.quantizer = quantizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "del quantizer_gpu "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "bento_stylesheets": {
   "bento/extensions/flow/main.css": true,
   "bento/extensions/kernel_selector/main.css": true,
   "bento/extensions/kernel_ui/main.css": true,
   "bento/extensions/new_kernel/main.css": true,
   "bento/extensions/system_usage/main.css": true,
   "bento/extensions/theme/main.css": true
  },
  "disseminate_notebook_id": {
   "notebook_id": "317645692710463"
  },
  "disseminate_notebook_info": {
   "bento_version": "20200809-210313",
   "description": "",
   "hide_code": false,
   "hipster_group": "",
   "kernel_build_info": {
    "deps": [
     "//deeplearning/projects/faiss:pyfaiss_gpu",
     "//deeplearning/projects/faiss/manifold:faiss_manifold_py",
     "//memcache/client/py:CacheClientWrapper"
    ],
    "external_deps": []
   },
   "no_uii": true,
   "notebook_number": "324202",
   "others_can_edit": false,
   "reviewers": "",
   "revision_id": "2716887035302012",
   "tags": "",
   "tasks": "",
   "title": "assign_on_gpu"
  },
  "kernelspec": {
   "display_name": "faiss",
   "language": "python",
   "name": "bento_kernel_faiss"
  },
  "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.7.5+"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}