adanns/adanns-ivf-unoptimized.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "04ace433",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import faiss\n",
    "import time\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import csv\n",
    "from os import path, makedirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e8a1adec",
   "metadata": {},
   "outputs": [],
   "source": [
    "root = 'path/to/embeddings/'\n",
    "\n",
    "D = 2048\n",
    "D_rr_search = 16 # to load high-D database and queryset for AR with rr models\n",
    "\n",
    "method = 'adanns' # adanns, mg-ivf-rr, mg-ivf-svd\n",
    "dataset = '1K' # 1K, 4K, V2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f98c143d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data_helper(config, D_load=2048):\n",
    "    db_csv = dataset + '_train_' + config + '-X.npy'\n",
    "    query_csv = dataset + '_val_' + config + '-X.npy'\n",
    "    db_label_csv = dataset + '_train_' + config + '-y.npy'\n",
    "    query_label_csv = dataset + '_val_' + config + '-y.npy'\n",
    "    \n",
    "    if dataset == 'V2':\n",
    "        db_csv = \"1K_train_\" + config + '-X.npy'\n",
    "        db_label_csv = \"1K_train_\" + config + '-y.npy'\n",
    "\n",
    "    db_load = np.ascontiguousarray(np.load(root+db_csv)[:, :D_load], dtype=np.float32)\n",
    "    qy_load = np.ascontiguousarray(np.load(root+query_csv)[:, :D_load], dtype=np.float32)\n",
    "    db_labels = np.load(root+db_label_csv)\n",
    "    query_labels = np.load(root+query_label_csv)\n",
    "\n",
    "    faiss.normalize_L2(db_load)\n",
    "    faiss.normalize_L2(qy_load)\n",
    "\n",
    "    return db_load, qy_load, db_labels, query_labels\n",
    "\n",
    "\n",
    "def load_construct_data(D_construct, D_rr_svd, ncentroids):\n",
    "    if method == 'adanns':\n",
    "        config = f'mrl1_e0_ff{D_construct}'\n",
    "    elif method == 'mg-ivf-rr':\n",
    "        config = f'mrl0_e0_ff{D_construct}'\n",
    "    elif method == 'mg-ivf-svd':\n",
    "        config = f'mrl0_e0_rr{D_construct}_svd{D_rr_svd}'\n",
    "    else:\n",
    "        raise Exception(\"Unsupported ANNS method.\")\n",
    "    db_construct, qy_construct, db_labels, query_labels = load_data_helper(config, D_construct)\n",
    "        \n",
    "    print(\"Cluster Contruction DB: \", db_construct.shape)\n",
    "    print(\"Cluster Construction queries:\", qy_construct.shape)\n",
    "    \n",
    "    # Load kmeans index and centroids with shape (centroid, D_construct)\n",
    "    size = str(ncentroids)+'ncentroid_'+str(D_construct)+'Dc'\n",
    "    if dataset == 'V2': # V2 is only a test set, change to 1K\n",
    "        dataset = '1K'\n",
    "    index_file = root+'index_files/'+method+dataset+'_kmeans_'+size\n",
    "\n",
    "    centroids_path = root+'kmeans/'+method+'ncentroids'+str(ncentroids)+\"_\"+str(D_construct)+'Dc_'+dataset+'.npy'\n",
    "    centroids = np.load(centroids_path)\n",
    "    print(\"Loaded centroids: \", centroids.shape, centroids_path)\n",
    "    \n",
    "    return db_construct, qy_construct, db_labels, query_labels, centroids, index_file\n",
    "\n",
    "def load_search_data(D_search):\n",
    "    if method == 'adanns':\n",
    "        config = f'mrl1_e0_ff{D_search}'\n",
    "    elif method in ['mg-ivf-rr', 'mg-ivf-svd']:\n",
    "        config = f'mrl0_e0_ff{D_search}'\n",
    "    else:\n",
    "        raise Exception(\"Unsupported ANNS method.\")\n",
    "    db_search, qy_search, _ , _ = load_data_helper(config, D_search)\n",
    "\n",
    "    return db_search, qy_search\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "55ff5ca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_cluster(val_classes, db_classes, neighbors, k): \n",
    "    APs, topk, recall = [], [], []\n",
    "    cluster_size = neighbors.shape[0]\n",
    "    for i in range(cluster_size):\n",
    "        target = val_classes[i]\n",
    "        indices = neighbors[i][:k] # k neighbor list for ith val vector\n",
    "        labels = db_classes[indices]\n",
    "        matches = (labels == target)\n",
    "        \n",
    "        # topk\n",
    "        hits = np.sum(matches)\n",
    "        if hits>0:\n",
    "            topk.append(1)\n",
    "        else:\n",
    "            topk.append(0)\n",
    "        \n",
    "        # recall\n",
    "        recall.append(np.sum(matches)/1300)\n",
    "        \n",
    "        # precision values\n",
    "        tps = np.cumsum(matches)\n",
    "        precs = tps.astype(float) / np.arange(1, k + 1, 1)\n",
    "        APs.append(np.sum(precs[matches.squeeze()]) / k)\n",
    "        \n",
    "    return np.mean(recall), np.mean(topk), np.mean(APs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5a66d9b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_closest_centroids(queries, centroids, D_shortlist):\n",
    "    centroid_index = faiss.IndexFlatL2(D_shortlist)\n",
    "    xq_shortlist = np.ascontiguousarray(queries[:, :D_shortlist], dtype=np.float32)\n",
    "    xc_shortlist = np.ascontiguousarray(centroids[:, :D_shortlist], dtype=np.float32)\n",
    "    faiss.normalize_L2(xq_shortlist)\n",
    "    faiss.normalize_L2(xc_shortlist)\n",
    "    \n",
    "    centroid_index.add(xc_shortlist)\n",
    "    _, I = centroid_index.search(xq_shortlist, 1)\n",
    "\n",
    "    return I"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b69174c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['d_c', 'd_s', 'd_shortlist', 'ncentroid', 'top1', 'recall@100', 'mAP@100', 'overlap']\n",
      "Cluster Contruction DB:  (1281167, 2048)\n",
      "Cluster Construction queries: (10000, 2048)\n",
      "\n",
      "Loaded kmeans index: 1K_kmeans_1024ncentroid_2048d\n",
      "Loaded centroids:  (1024, 2048) ../../inference_array/resnet50/kmeans/mrl/ncentroids1024_2048d_1K.npy\n",
      "Linear scan with d =  8\n",
      "[2048, 8, 2048, 1024, 0.5351, 0.05091163624126155, 0.605462860936279, 0]\n",
      "d_c:2048, d_s: 8, ncentroid: 1024\n",
      "Recall@100:  0.05091163624126155\n",
      "Top1:  0.5351\n",
      "Linear scan with d =  16\n",
      "[2048, 32, 2048, 1024, 0.5732, 0.051922105632729296, 0.6227463073287828, 0]\n",
      "d_c:2048, d_s: 32, ncentroid: 1024\n",
      "Recall@100:  0.051922105632729296\n",
      "Top1:  0.5732\n",
      "Linear scan with d =  64\n",
      "[2048, 64, 2048, 1024, 0.5785, 0.05198921759119408, 0.6243867377357402, 0]\n",
      "d_c:2048, d_s: 64, ncentroid: 1024\n",
      "Recall@100:  0.05198921759119408\n",
      "Top1:  0.5785\n",
      "Linear scan with d =  128\n",
      "[2048, 128, 2048, 1024, 0.5802, 0.05199322657824928, 0.6247377244191619, 0]\n",
      "d_c:2048, d_s: 128, ncentroid: 1024\n",
      "Recall@100:  0.05199322657824928\n",
      "Top1:  0.5802\n",
      "Linear scan with d =  256\n",
      "[2048, 256, 2048, 1024, 0.5801, 0.05199830116619949, 0.6250036507917335, 0]\n",
      "d_c:2048, d_s: 256, ncentroid: 1024\n",
      "Recall@100:  0.05199830116619949\n",
      "Top1:  0.5801\n",
      "Linear scan with d =  512\n",
      "[2048, 512, 2048, 1024, 0.5803, 0.05199979769465816, 0.6251013861127136, 0]\n",
      "d_c:2048, d_s: 512, ncentroid: 1024\n",
      "Recall@100:  0.05199979769465816\n",
      "Top1:  0.5803\n",
      "Linear scan with d =  1024\n",
      "[2048, 1024, 2048, 1024, 0.5766, 0.05198958841622709, 0.6249844689120537, 0]\n",
      "d_c:2048, d_s: 1024, ncentroid: 1024\n",
      "Recall@100:  0.05198958841622709\n",
      "Top1:  0.5766\n",
      "Total Time for 8 configs = 211.981683\n"
     ]
    }
   ],
   "source": [
    "D_search_list = [8, 16, 32, 64, 128, 256, 512, 1024]\n",
    "ncentroids = 1024\n",
    "\n",
    "for D in [2048]:\n",
    "    k=100\n",
    "    D_rr_svd = D\n",
    "    D_construct_list = [D]\n",
    "    D_shortlist_list = [D]\n",
    "\n",
    "    header = ['d_construct', 'd_search', 'd_shortlist', 'ncentroid', 'top1', 'recall@'+str(k), 'mAP@'+str(k)]\n",
    "    print(header)\n",
    "    with open('kmeans_metrics.csv', 'w', encoding='UTF8', newline='') as f:\n",
    "        writer = csv.writer(f)\n",
    "        writer.writerow(header)\n",
    "\n",
    "    start = time.time()\n",
    "    for D_c in D_construct_list:\n",
    "        # Load all construction data (database, queries, centroids)\n",
    "        xb_construct, xq_construct, db_labels, query_labels, centroids, index_file = load_construct_data(D_c, D_rr_svd, ncentroids)\n",
    "\n",
    "        cpu_index = faiss.read_index(index_file+'.index')\n",
    "        index = faiss.index_cpu_to_all_gpus(cpu_index)\n",
    "        print(\"\\nLoaded kmeans index:\", index_file.split(\"/\")[-1])\n",
    "\n",
    "        # construct lookup table of centroid --> vectors, i.e. inverted lists\n",
    "        _, I_db = index.search(xb_construct, 1)\n",
    "        lut_db = {}\n",
    "        for c in np.unique(I_db):\n",
    "            lut_db[c] = np.argwhere(I_db==c)[:,0]\n",
    "\n",
    "        for D_search in D_search_list:\n",
    "            print(\"Linear scan with D_s = \", D_search)\n",
    "            xb_search, xq_search = load_search_data(D_search)\n",
    "\n",
    "            for D_shortlist in D_shortlist_list:\n",
    "                # Currently, D_shortlist <= D_search is supported as we slice centroids for adanns\n",
    "                I_q = get_closest_centroids(xq_search, centroids, D_shortlist)\n",
    "                lut_q = {}\n",
    "\n",
    "                start = time.time()\n",
    "                recall, topk, mAP = [], [], []\n",
    "\n",
    "                #Iterate over all centroids assigned to each\n",
    "                for c in np.unique(I_q):\n",
    "                    lut_q[c] = np.argwhere(I_q==c)[:,0]\n",
    "                    exact_cpu_index = faiss.IndexFlatL2(D_search)\n",
    "\n",
    "                    # add cluster vectors to index and search only queries that map to that cluster\n",
    "                    exact = faiss.index_cpu_to_all_gpus(exact_cpu_index)\n",
    "                    cluster_db = np.ascontiguousarray(xb_search[lut_db[c]][:, :D_search], np.float32)\n",
    "                    cluster_query = np.ascontiguousarray(xq_search[lut_q[c]][:, :D_search], np.float32)\n",
    "                    faiss.normalize_L2(cluster_db)\n",
    "                    faiss.normalize_L2(cluster_query)\n",
    "                    exact.add(cluster_db)\n",
    "                    Dist, Ind = exact.search(cluster_query, k)\n",
    "\n",
    "                    # replace cluster-specific indices with original database indices for eval\n",
    "                    cluster_db_labels = db_labels[lut_db[c]]\n",
    "                    cluster_query_labels = query_labels[lut_q[c]]\n",
    "\n",
    "                    nn_1 = Ind[:, 0]\n",
    "                    pred_1 = cluster_db_labels[nn_1]\n",
    "                    hits = np.sum(pred_1 == cluster_query_labels)\n",
    "                    topk.append(hits)\n",
    "\n",
    "                    rl, tk, mp = eval_cluster(cluster_query_labels, cluster_db_labels, Ind, k)\n",
    "                    recall.append(rl)\n",
    "                    mAP.append(mp)\n",
    "                row = [D_c, D_search, D_shortlist, ncentroids, np.sum(topk)/xq_search.shape[0], np.mean(recall), np.mean(mAP)]\n",
    "                print(row)\n",
    "\n",
    "                with open('kmeans_metrics.csv', 'a', encoding='UTF8', newline='') as f:\n",
    "                    writer = csv.writer(f)\n",
    "                    writer.writerow(row)\n",
    "                print(\"d_c:%d, d_s: %d, ncentroid: %d\" %(D_c, D_search, ncentroids))\n",
    "                print(\"Recall@100: \", np.mean(recall))\n",
    "                print(\"Top1: \", np.sum(topk)/xq_search.shape[0])\n",
    "\n",
    "    print(\"Total Time for %d configs = %f\" % (len(D_search_list) * len(ncentroids) * len(D_construct_list), time.time() - start))"
   ]
  }
 ],
 "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.11.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "5c7b89af1651d0b8571dde13640ecdccf7d5a6204171d6ab33e7c296e100e08a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}