-
-
-
-
-
-{% endblock %}
+{% extends "general.html" %}
+{% block content %}
+
+
Billion-Scale Approximate Nearest Neighbor Search Challenge: NeurIPS'21 competition track
+
+
+
+
+
+{% endblock %}
diff --git a/tests/recall_tests.py b/tests/recall_tests.py
deleted file mode 100755
index 95099a66d..000000000
--- a/tests/recall_tests.py
+++ /dev/null
@@ -1,243 +0,0 @@
-import numpy as np
-import sys
-
-from benchmark.plotting.metrics import get_recall_values
-from benchmark.datasets import DATASETS
-
-ASSERT= True # Stop unit tests on first failure
-
-GT_MIN_SIZE = 20 # Require ground truth with at least this length for each query
-
-def main_tests():
- #
- # test recall computation on fake responses
- #
-
- def test_recall( true_ids, true_dists, run_ids, count, expected_no_ties, expected_with_ties ):
- '''This function will test the two forms of recall (with and without considering ties.)'''
-
- # compute recall, don't consider ties
- recall = get_recall_values( (true_ids, true_dists), run_ids, count, False)
- expected = 1.0
- print("compute recall(don't consider ties)=%f" % recall[0], "expected recall=%f" % expected_no_ties)
- if ASSERT:
- assert recall[0]==expected_no_ties
- print("passed")
-
- # compute recall, consider ties
- recall = get_recall_values( (true_ids, true_dists), run_ids, count, True)
- expected = 1.0
- print("compute recall(consider ties)=%f num_queries_with_ties=%d" % (recall[0], recall[3]), "expected recall=%f" % expected_with_ties)
- if ASSERT:
- assert recall[0]==expected_with_ties
- print("passed")
-
- print()
-
- print("TEST: fake query response with no distance ties, 1 query and k=3")
- true_ids = np.array([ [ 0, 1, 2 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 2 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, 1.0, 1.0 )
-
- print("TEST: fake query response with no distance ties but not 1.0 recall, 1 query and k=3")
- true_ids = np.array([ [ 0, 1, 2 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, 2.0/3.0, 2.0/3.0 )
-
- print("TEST: fake query response with no ties, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2 ], [ 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0 ], [ 0.0, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 2 ], [ 2, 1, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, 1.0, 1.0 )
-
- print("TEST: fake query response with no distance ties, 1 query, k=3, GT array is larger than run array")
- true_ids = np.array([ [ 0, 1, 2, 3 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0, 3.0 ] ])
- run_ids = np.array([ [ 0, 1, 2 ] ])
- count=3
- print("yuk true_ids=", true_ids.shape, "run_ids=", run_ids.shape)
- test_recall( true_ids, true_dists, run_ids, count, 1.0, 1.0 )
-
- print("TEST: fake query response with an out-of-bounds distance ties, 1 query, k=3, GT array is larger than run array.")
- true_ids = np.array([ [ 0, 1, 2, 3 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 2 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, 1.0, 1.0 )
-
- # this is from bigann GT and query set. The GT arrays are size=11 but run array is 10 and there are no ties to consider
- print("TEST: from bigann-1B...")
- true_ids = np.array([ [937541801, 221456167, 336118969, 971823307, 267986685, 544978851, 815975675, 615142927, 640142873, 994367459, 504814] ] )
- true_dists = np.array([ [55214., 58224., 58379., 58806., 59251., 59256., 60302., 60573., 60843., 60950., 61125.] ] )
- run_ids = np.array([ [221456167, 336118969, 971823307, 640142873, 994367459, 504814, 87356234, 628179290, 928121617, 397551598 ] ] )
- count=10
- test_recall( true_ids, true_dists, run_ids, count, 0.5, 0.5 )
-
- print("TEST: fake query response with ties at beginning, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 0.0, 1.0, 2.0 ], [ 0.0, 0.0, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0,2.0])/count, np.mean([2.0,2.0])/count)
-
- print("TEST: fake query response with ties at beginning and ties have small diff, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 0.0+1e-6-1e-7, 1.0, 2.0 ], [ 0.0, 0.0+1e-6-1e-7, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0,2.0])/count, np.mean([2.0,2.0])/count)
-
- print("TEST: fake query response with possible ties at beginning but diff is just beyond the 1e-6 threshold, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 0.0+1e-6, 1.0, 2.0 ], [ 0.0, 0.0+1e-6, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0,2.0])/count, np.mean([2.0,2.0])/count )
-
- print("TEST: fake query response with ties in middle, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 1.0, 2.0 ], [ 0.0, 1.0, 1.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0,2.0])/count, np.mean([2.0,2.0])/count)
-
- print("TEST: fake query response with ties at count-1 and 1 tie after, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0, 2.0 ], [ 0.0, 1.0, 2.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0, 2.0])/float(count), np.mean([3.0, 3.0])/float(count) )
-
- print("TEST: fake query response with ties at count-1 and 1 tie after and 1 after that that is cloe, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3 ], [ 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0, 2.0 ], [ 0.0, 1.0, 2.0, 2.0 ] ])
- run_ids = np.array([ [ 0, 1, 3 ], [ 3, 2, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0, 2.0])/float(count), np.mean([3.0, 3.0])/float(count) )
-
- print("TEST: fake query response with ties at count-1 and several close ties after, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3, 4, 5 ], [ 5, 4, 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 1.0, 2.0, 2.0, 2.0+1e-6-1e-7, 2.0+1e-6 ], [ 0.0, 1.0, 2.0, 2.0, 2.0+1e-6-1e-7, 2.0+1e-6 ] ])
- run_ids = np.array([ [ 0, 1, 4 ], [ 5, 4, 0 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0,2.0])/count, np.mean([3.0, 2.0])/float(count) )
-
- print("TEST: fake query response with two independent tie groups, 2 queries and k=3")
- true_ids = np.array([ [ 0, 1, 2, 3, 4, 5 ], [ 5, 4, 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 0.0, 2.0, 3.0, 3.0, 4.0 ], [ 0.0, 0.0, 2.0, 3.0, 3.0, 4.0 ] ])
- run_ids = np.array([ [ 0, 1, 5 ], [ 5, 4, 1 ] ])
- count=3
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0, 2.0])/float(count), np.mean([2.0, 2.0])/float(count) )
-
- print("TEST: fake query response with two independent tie groups, 2 queries and k=4")
- true_ids = np.array([ [ 0, 1, 2, 3, 4, 5, 6 ], [ 6, 5, 4, 3, 2, 1, 0 ] ])
- true_dists = np.array([ [ 0.0, 0.0, 2.0, 3.0, 3.0, 4.0, 5.0 ], [ 0.0, 0.0, 2.0, 3.0, 3.0, 4.0, 5.0 ] ])
- run_ids = np.array([ [ 0, 1, 5, 7 ], [ 5, 4, 1, 7 ] ])
- count=4
- test_recall( true_ids, true_dists, run_ids, count, np.mean([2.0, 2.0])/float(count), np.mean([2.0, 2.0])/float(count) )
-
- #
- # dataset tests
- #
- def test_GT_monotonicity( dset, increasing=True ):
- print("TEST: %s, checking GT distances monotonicity" % dset)
- dataset = DATASETS[dset]()
- gt = dataset.get_groundtruth()
- if ASSERT: assert len(gt)==2
- true_ids = gt[0]
- true_dists = gt[1]
- if ASSERT:
- assert true_ids.shape[1]==true_dists.shape[1]
- assert true_ids.shape[1]>=GT_MIN_SIZE
- assert true_dists.shape[1]>=GT_MIN_SIZE
- func = monotone_increasing if increasing else monotone_decreasing
- for i in range(true_dists.shape[0]):
- mtest = func(true_dists[i])
- if ASSERT: assert mtest==True
- print()
-
- print("TEST: sanity check the monotone functions")
- mtest = monotone_increasing([0,1,2,3,4,5])
- if ASSERT: assert mtest==True
- mtest = monotone_increasing([0,0,0,3,4,5])
- if ASSERT: assert mtest==True
- mtest = monotone_increasing([3,4,5,4,3,4,5])
- if ASSERT: assert mtest==False
- mtest = monotone_increasing([5,4,4,3,2,1])
- if ASSERT: assert mtest==False
- print()
-
- # check GT dist increasing monotonicity for each knn dataset
- test_GT_monotonicity( "bigann-1B" )
- test_GT_monotonicity( "deep-1B" )
- test_GT_monotonicity( "msturing-1B" )
- test_GT_monotonicity( "msspacev-1B" )
- test_GT_monotonicity( "text2image-1B", increasing=False)
-
- #
- # test recall on actual datasets
- #
- def extract_GT_monotonicity( dset, row, c1, c2):
- print("TEST: %s, extraction" % dset, row, c1, c2)
- dataset = DATASETS[dset]()
- gt = dataset.get_groundtruth()
- true_dists = gt[1]
- lst = true_dists[row,c1:c2]
- print(lst)
- mtest = monotone_increasing(lst)
- print(mtest)
- print()
-
- def test_GT_as_query( dset, count ):
- print("TEST: %s, using GT as query, k=10" % dset)
- dataset = DATASETS[dset]()
- gt = dataset.get_groundtruth()
- if ASSERT: assert len(gt)==2
- true_ids = gt[0]
- true_dists = gt[1]
- if ASSERT:
- assert true_ids.shape[1]==true_dists.shape[1]
- assert true_ids.shape[1]>=GT_MIN_SIZE
- assert true_dists.shape[1]>=GT_MIN_SIZE
- run_ids = np.copy( gt[0] )[:,0:count] # create a query set from GT truncated at k
- test_recall( true_ids, true_dists, run_ids, count, 1.0, 1.0 )
-
- # test GT as query for each dataset
- test_GT_as_query( "bigann-1B", 10 )
- test_GT_as_query( "deep-1B", 10 )
- test_GT_as_query( "text2image-1B", 10 )
- test_GT_as_query( "msturing-1B", 10 )
- test_GT_as_query( "msspacev-1B", 10 )
-
- sys.exit(0)
-
-
-
-#
-# useful functions
-#
-import itertools
-import operator
-
-def monotone_increasing(lst):
- pairs = zip(lst, lst[1:])
- bools = list(itertools.starmap(operator.le, pairs))
- #print(type(bools), len(bools), bools)
- return all( bools )
-
-def monotone_decreasing(lst):
- pairs = zip(lst, lst[1:])
- bools = list(itertools.starmap(operator.ge, pairs))
- #print(type(lst), lst)
- return all( bools )
-
-def monotone(lst):
- return monotone_increasing(lst) or monotone_decreasing(lst)
-
-if __name__ == "__main__":
- main_tests()
diff --git a/tests/tests.sh b/tests/tests.sh
deleted file mode 100755
index 2e59a08d9..000000000
--- a/tests/tests.sh
+++ /dev/null
@@ -1,6 +0,0 @@
-#!/bin/bash
-
-# You should run this script from the repo top-level directory
-
-PYTHONPATH="." python tests/recall_tests.py
-
diff --git a/track1_baseline_faiss/README.md b/track1_baseline_faiss/README.md
deleted file mode 100644
index 090f80ed2..000000000
--- a/track1_baseline_faiss/README.md
+++ /dev/null
@@ -1,170 +0,0 @@
-# Running the Faiss baselines
-
-## Installing software
-
-In addition to this repository, running the baseline code requires a conda install with Faiss
-
-```bash
-wget https://repo.anaconda.com/archive/Anaconda3-2020.11-Linux-x86_64.sh
-
-bash Anaconda3-2020.11-Linux-x86_64.sh
-
-# follow instructions and run profile.sh to get a working conda
-
-conda create -n faiss_1.7.1 python=3.8
-conda activate faiss_1.7.1
-conda install -c pytorch faiss-cpu
-```
-
-All instructions below are supposed to be run from the root of the repository.
-To make the package accessible, set `export PYTHONPATH=.`
-
-## Downloading the data
-
-To download the data (database files, query files and ground truth, do
-```
-mkdir data/ # this is where all the data goes, a symlink is fine
-python track1_baseline_faiss/baseline_faiss.py --dataset deep-1B --prepare
-```
-The available datasets are bigann-1B deep-1B ssnpp-1B text2image-1B msturing-1B msspacev-1B.
-To download the largest files, `--prepare` will use axel or azcopy. Make sure that they are in the path.
-
-Replace the -1B suffix with -100M or -10M to get a subset of each dataset (only the relevant fraction of the database will be downloaded).
-This is useful for small-scale experiments.
-
-## Building the index
-
-There are several types of indexes in Faiss.
-Here we focus on IVF variants with PQ compression as recommended [here](https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index#if-100m---1b-ivf1048576_hnsw32) and evaluated [here](https://github.com/facebookresearch/faiss/wiki/Indexing-1G-vectors#1b-datasets).
-
-The problem is that they require very large codebooks to define the IVF clusters.
-This is fine (kind of) when a GPU is available to run the clustering, but not on CPU only.
-Therefore, we perform a two-level clustering with n' = sqrt(ncentroids) first level cluster and n' clusterings of size n' at a refined level.
-Then all n' * n' sub-clusters are indexed together in an IVF_HNSW.
-
-This writes like:
-
-```bash
-python -u track1_baseline_faiss/baseline_faiss.py --dataset deep-1B \
- --indexkey OPQ64_128,IVF1048576_HNSW32,PQ64x4fsr \
- --maxtrain 100000000 \
- --two_level_clustering \
- --build \
- --add_splits 30 \
- --indexfile data/track1_baseline_faiss/deep-1B.IVF1M_2level_PQ64x4fsr.faissindex \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80
-```
-
-This works for deep-1B bigann-1B msturing-1B msspacev-1B.
-
-For ssnpp-1B, the type of index has to be adjusted a bit because the Faiss PQ64x4fsr does not support range search (see [the documentation of Faiss index types](https://github.com/facebookresearch/faiss/wiki/The-index-factory#encodings) for an explanation of the difference).
-
-Therefore, we use a slightly slower index type: PQ32. This gives:
-```bash
-python -u track1_baseline_faiss/baseline_faiss.py --dataset ssnpp-1B \
- --indexkey OPQ64_128,IVF1048576_HNSW32,PQ32 \
- --maxtrain 100000000 \
- --two_level_clustering \
- --build \
- --add_splits 30 \
- --indexfile data/track1_baseline_faiss/ssnpp-1B.IVF1M_2level_PQ23.faissindex \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80
-```
-
-The results on text2image-1B with the same index types are a lot worse.
-This is probably due to the very lossy PQ compression.
-
-## Running the evaluation
-
-### Getting the pre-built indexes
-
-Pre-built indexes are available.
-To download them
-
-```bash
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/deep-1B.IVF1M_2level_PQ64x4fsr.faissindex -P data/
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/bigann-1B.IVF1M_2level_PQ64x4fsr.faissindex -P data/
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/msturing-1B.IVF1M_2level_PQ64x4fsr.faissindex -P data/
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/msspacev-1B.IVF1M_2level_PQ64x4fsr.faissindex -P data/
-
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/ssnpp-1B.IVF1M_2level_PQ32.faissindex -P data/
-wget https://dl.fbaipublicfiles.com/billion-scale-ann-benchmarks/track1_baseline_faiss/text2image-1B.IVF1M_2level_PQ32.faissindex -P data/
-
-```
-
-
-### Running the evaluation
-
-
-
-
-The evaluation proceeds by loading the index and looping over a set of search-time parameters that obtain different speed-accuracy tradeoffs.
-
-This writes as:
-```bash
-
-params="
-nprobe=1,quantizer_efSearch=4
-nprobe=2,quantizer_efSearch=4
-...
-nprobe=512,quantizer_efSearch=256
-nprobe=512,quantizer_efSearch=512
-nprobe=1024,quantizer_efSearch=512
-"
-
-python track1_baseline_faiss/baseline_faiss.py \
- --dataset deep-1B --indexfile data/deep-1B.IVF1M_2level_PQ64x4fsr.faissindex \
- --search --searchparams $params
-
-```
-
-The sets of parameters per dataset are listed in [this GIST](https://gist.github.com/mdouze/bb71032f0b3bf3cc9bdaa6ff1287c144).
-They are ordered from fastest / least accurate to slowest / most accurate.
-
-### Results
-
-The results should look like:
-
-```
-parameters inter@ 10 time(ms/q) nb distances #runs
-nprobe=1,quantizer_efSearch=4 0.1738 0.00327 12210374 92
-nprobe=2,quantizer_efSearch=4 0.2394 0.00424 24328050 71
-nprobe=2,quantizer_efSearch=8 0.2879 0.00545 24278048 56
-...
-nprobe=512,quantizer_efSearch=256 0.6877 0.75883 5896044691 1
-nprobe=512,quantizer_efSearch=512 0.6886 0.77421 5890639041 1
-nprobe=1024,quantizer_efSearch=512 0.6886 1.46841 11607413418 1
-```
-
-This means that by setting the parameters `nprobe=2,quantizer_efSearch=4`, we obtain 0.2394 recall @ 10 (aka inter @10) for that dataset, the search will take 0.00327 ms per query (305810 QPS).
-The total number of distances computed for all queries is 24328050 and this measurement was obtained in 71 runs (to reduce jitter in time measurements).
-
-
-### Plots
-
-The speed-accuracy tradeoff plots are here (with 32 threads on a given 2.2Ghz machine):
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-### Determining the optimal search-time parameters
-
-The Pareto-optimal parameter combinations can be obtained by a random exploration of the parameter space, as described [here](https://github.com/facebookresearch/faiss/wiki/Index-IO,-cloning-and-hyper-parameter-tuning#auto-tuning-the-runtime-parameters).
-To perform this operation, do:
-```bash
-python track1_baseline_faiss/baseline_faiss.py \
- --dataset deep-1B --indexfile data/deep-1B.IVF1M_2level_PQ64x4fsr.faissindex \
- --search
-```
diff --git a/track1_baseline_faiss/__init__.py b/track1_baseline_faiss/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/track1_baseline_faiss/baseline_faiss.py b/track1_baseline_faiss/baseline_faiss.py
deleted file mode 100644
index 4446401e2..000000000
--- a/track1_baseline_faiss/baseline_faiss.py
+++ /dev/null
@@ -1,791 +0,0 @@
-import os
-import sys
-import time
-import pdb
-import gc
-import numpy as np
-import faiss
-import argparse
-import resource
-
-import benchmark.datasets
-from benchmark.datasets import DATASETS
-from benchmark.plotting import eval_range_search
-
-####################################################################
-# Index building functions
-####################################################################
-
-
-def two_level_clustering(xt, nc1, nc2, clustering_niter=25, spherical=False):
- d = xt.shape[1]
-
- print(f"2-level clustering of {xt.shape} nb clusters = {nc1}*{nc2} = {nc1*nc2}")
- print("perform coarse training")
-
- km = faiss.Kmeans(
- d, nc1, verbose=True, niter=clustering_niter,
- max_points_per_centroid=2000,
- spherical=spherical
- )
- km.train(xt)
-
- print()
-
- # coarse centroids
- centroids1 = km.centroids
-
- print("assigning the training set")
- t0 = time.time()
- _, assign1 = km.assign(xt)
- bc = np.bincount(assign1, minlength=nc1)
- print(f"done in {time.time() - t0:.2f} s. Sizes of clusters {min(bc)}-{max(bc)}")
- o = assign1.argsort()
- del km
-
- # train sub-clusters
- i0 = 0
- c2 = []
- t0 = time.time()
- for c1 in range(nc1):
- print(f"[{time.time() - t0:.2f} s] training sub-cluster {c1}/{nc1}\r", end="", flush=True)
- i1 = i0 + bc[c1]
- subset = o[i0:i1]
- assert np.all(assign1[subset] == c1)
- km = faiss.Kmeans(d, nc2, spherical=spherical)
- xtsub = xt[subset]
- km.train(xtsub)
- c2.append(km.centroids)
- i0 = i1
- print(f"done in {time.time() - t0:.2f} s")
- return np.vstack(c2)
-
-
-def unwind_index_ivf(index):
- if isinstance(index, faiss.IndexPreTransform):
- assert index.chain.size() == 1
- vt = faiss.downcast_VectorTransform(index.chain.at(0))
- index_ivf, vt2 = unwind_index_ivf(faiss.downcast_index(index.index))
- assert vt2 is None
- return index_ivf, vt
- if hasattr(faiss, "IndexRefine") and isinstance(index, faiss.IndexRefine):
- return unwind_index_ivf(faiss.downcast_index(index.base_index))
- if isinstance(index, faiss.IndexIVF):
- return index, None
- else:
- return None, None
-
-
-def build_index(args, ds):
- nq, d = ds.nq, ds.d
- nb, d = ds.nq, ds.d
-
- if args.buildthreads == -1:
- print("Build-time number of threads:", faiss.omp_get_max_threads())
- else:
- print("Set build-time number of threads:", args.buildthreads)
- faiss.omp_set_num_threads(args.buildthreads)
-
- metric_type = (
- faiss.METRIC_L2 if ds.distance() == "euclidean" else
- faiss.METRIC_INNER_PRODUCT if ds.distance() in ("ip", "angular") else
- 1/0
- )
- print("metric type", metric_type)
- index = faiss.index_factory(d, args.indexkey, metric_type)
-
- index_ivf, vec_transform = unwind_index_ivf(index)
- if vec_transform is None:
- vec_transform = lambda x: x
- else:
- vec_transform = faiss.downcast_VectorTransform(vec_transform)
-
- if args.by_residual != -1:
- by_residual = args.by_residual == 1
- print("setting by_residual = ", by_residual)
- index_ivf.by_residual # check if field exists
- index_ivf.by_residual = by_residual
-
- if index_ivf:
- print("Update add-time parameters")
- # adjust default parameters used at add time for quantizers
- # because otherwise the assignment is inaccurate
- quantizer = faiss.downcast_index(index_ivf.quantizer)
- if isinstance(quantizer, faiss.IndexRefine):
- print(" update quantizer k_factor=", quantizer.k_factor, end=" -> ")
- quantizer.k_factor = 32 if index_ivf.nlist < 1e6 else 64
- print(quantizer.k_factor)
- base_index = faiss.downcast_index(quantizer.base_index)
- if isinstance(base_index, faiss.IndexIVF):
- print(" update quantizer nprobe=", base_index.nprobe, end=" -> ")
- base_index.nprobe = (
- 16 if base_index.nlist < 1e5 else
- 32 if base_index.nlist < 4e6 else
- 64)
- print(base_index.nprobe)
- elif isinstance(quantizer, faiss.IndexHNSW):
- print(" update quantizer efSearch=", quantizer.hnsw.efSearch, end=" -> ")
- if args.quantizer_add_efSearch > 0:
- quantizer.hnsw.efSearch = args.quantizer_add_efSearch
- else:
- quantizer.hnsw.efSearch = 40 if index_ivf.nlist < 4e6 else 64
- print(quantizer.hnsw.efSearch)
- if args.quantizer_efConstruction != -1:
- print(" update quantizer efConstruction=", quantizer.hnsw.efConstruction, end=" -> ")
- quantizer.hnsw.efConstruction = args.quantizer_efConstruction
- print(quantizer.hnsw.efConstruction)
-
-
- index.verbose = True
- if index_ivf:
- index_ivf.verbose = True
- index_ivf.quantizer.verbose = True
- index_ivf.cp.verbose = True
-
-
- maxtrain = args.maxtrain
- if maxtrain == 0:
- if 'IMI' in args.indexkey:
- maxtrain = int(256 * 2 ** (np.log2(index_ivf.nlist) / 2))
- elif index_ivf:
- maxtrain = 50 * index_ivf.nlist
- else:
- # just guess...
- maxtrain = 256 * 100
- maxtrain = max(maxtrain, 256 * 100)
- print("setting maxtrain to %d" % maxtrain)
-
- # train on dataset
- print(f"getting first {maxtrain} dataset vectors for training")
-
- xt2 = next(ds.get_dataset_iterator(bs=maxtrain))
-
- print("train, size", xt2.shape)
- assert np.all(np.isfinite(xt2))
-
- t0 = time.time()
-
- if (isinstance(vec_transform, faiss.OPQMatrix) and
- isinstance(index_ivf, faiss.IndexIVFPQFastScan)):
- print(" Forcing OPQ training PQ to PQ4")
- ref_pq = index_ivf.pq
- training_pq = faiss.ProductQuantizer(
- ref_pq.d, ref_pq.M, ref_pq.nbits
- )
- vec_transform.pq
- vec_transform.pq = training_pq
-
- if args.clustering_niter >= 0:
- print(("setting nb of clustering iterations to %d" %
- args.clustering_niter))
- index_ivf.cp.niter = args.clustering_niter
-
- train_index = None
- if args.train_on_gpu:
- print("add a training index on GPU")
- train_index = faiss.index_cpu_to_all_gpus(
- faiss.IndexFlatL2(index_ivf.d))
- index_ivf.clustering_index = train_index
-
- if args.two_level_clustering:
- sqrt_nlist = int(np.sqrt(index_ivf.nlist))
- assert sqrt_nlist ** 2 == index_ivf.nlist
-
- centroids_trainset = xt2
- if isinstance(vec_transform, faiss.VectorTransform):
- print(" training vector transform")
- vec_transform.train(xt2)
- print(" transform trainset")
- centroids_trainset = vec_transform.apply_py(centroids_trainset)
-
- centroids = two_level_clustering(
- centroids_trainset, sqrt_nlist, sqrt_nlist,
- spherical=(metric_type == faiss.METRIC_INNER_PRODUCT)
- )
-
- if not index_ivf.quantizer.is_trained:
- print(" training quantizer")
- index_ivf.quantizer.train(centroids)
-
- print(" add centroids to quantizer")
- index_ivf.quantizer.add(centroids)
-
- index.train(xt2)
- print(" Total train time %.3f s" % (time.time() - t0))
-
- if train_index is not None:
- del train_index
- index_ivf.clustering_index = None
- gc.collect()
-
- print("adding")
-
- t0 = time.time()
- if args.add_bs == -1:
- index.add(sanitize(ds.get_database()))
- else:
- i0 = 0
- nsplit = args.add_splits
- for sno in range(nsplit):
- print(f"============== SPLIT {sno}/{nsplit}")
- for xblock in ds.get_dataset_iterator(bs=args.add_bs, split=(nsplit, sno)):
- i1 = i0 + len(xblock)
- print(" adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
- i0, i1, ds.nb, time.time() - t0,
- faiss.get_mem_usage_kb()))
- index.add(xblock)
- i0 = i1
- gc.collect()
- if sno == args.stop_at_split:
- print("stopping at split", sno)
- break
-
- print(" add in %.3f s" % (time.time() - t0))
- if args.indexfile:
- print("storing", args.indexfile)
- faiss.write_index(index, args.indexfile)
-
- return index
-
-####################################################################
-# Evaluation functions
-####################################################################
-
-
-def compute_inter(a, b):
- nq, rank = a.shape
- ninter = sum(
- np.intersect1d(a[i, :rank], b[i, :rank]).size
- for i in range(nq)
- )
- return ninter / a.size
-
-def knn_search_batched(index, xq, k, bs):
- D, I = [], []
- for i0 in range(0, len(xq), bs):
- Di, Ii = index.search(xq[i0:i0 + bs], k)
- D.append(Di)
- I.append(Ii)
- return np.vstack(D), np.vstack(I)
-
-def eval_setting_knn(index, xq, gt, k=0, inter=False, min_time=3.0, query_bs=-1):
- nq = xq.shape[0]
- gt_I, gt_D = gt
-
- ivf_stats = faiss.cvar.indexIVF_stats
- ivf_stats.reset()
- nrun = 0
- t0 = time.time()
- while True:
- if query_bs == -1:
- D, I = index.search(xq, k)
- else:
- D, I = knn_search_batched(index, xq, k, query_bs)
- nrun += 1
- t1 = time.time()
- if t1 - t0 > min_time:
- break
- ms_per_query = ((t1 - t0) * 1000.0 / nq / nrun)
-
- if inter:
- rank = k
- inter_measure = compute_inter(gt_I[:, :rank], I[:, :rank])
- print("%.4f" % inter_measure, end=' ')
- else:
- for rank in 1, 10, 100:
- n_ok = (I[:, :rank] == gt_I[:, :1]).sum()
- print("%.4f" % (n_ok / float(nq)), end=' ')
- print(" %9.5f " % ms_per_query, end=' ')
-
- if ivf_stats.search_time == 0:
- # happens for IVFPQFastScan where the stats are not logged by default
- print("%12d %5.2f " % (ivf_stats.ndis / nrun, 0.0), end=' ')
- else:
- pc_quantizer = ivf_stats.quantization_time / ivf_stats.search_time * 100
- print("%12d %5.2f " % (ivf_stats.ndis / nrun, pc_quantizer), end=' ')
- print(nrun)
-
-def eval_setting_range(index, xq, gt, radius=0, inter=False, min_time=3.0, query_bs=-1):
- nq = xq.shape[0]
- gt_nres, gt_I, gt_D = gt
- gt_lims = np.zeros(nq + 1, dtype=int)
- gt_lims[1:] = np.cumsum(gt_nres)
- ivf_stats = faiss.cvar.indexIVF_stats
- ivf_stats.reset()
- nrun = 0
- t0 = time.time()
- while True:
- if query_bs == -1:
- lims, D, I = index.range_search(xq, radius)
- else:
- raise NotImplemented
- nrun += 1
- t1 = time.time()
- if t1 - t0 > min_time:
- break
- ms_per_query = ((t1 - t0) * 1000.0 / nq / nrun)
-
- ap = eval_range_search.compute_AP((gt_lims, gt_I, gt_D), (lims, I, D))
- print("%.4f" % ap, end=' ')
- print(" %9.5f " % ms_per_query, end=' ')
-
- print("%12d %5d " % (ivf_stats.ndis / nrun, D.size), end=' ')
- print(nrun)
-
-
-def result_header(ds, args):
-
- # setup the Criterion object
- if ds.search_type() == "range":
- header = (
- '%-40s AP time(ms/q) nb distances nb_res #runs' %
- "parameters"
- )
- crit = None
- elif args.inter:
- print("Optimize for intersection @ ", args.k)
- crit = faiss.IntersectionCriterion(ds.nq, args.k)
- header = (
- '%-40s inter@%3d time(ms/q) nb distances %%quantization #runs' %
- ("parameters", args.k)
- )
- else:
- print("Optimize for 1-recall @ 1")
- crit = faiss.OneRecallAtRCriterion(ds.nq, 1)
- header = (
- '%-40s R@1 R@10 R@100 time(ms/q) nb distances %%quantization #runs' %
- "parameters"
- )
- return header, crit
-
-def op_compute_bounds(ps, ops, cno):
- # lower_bound_t = 0.0
- # upper_bound_perf = 1.0
- bounds = np.array([0, 1], dtype="float64")
- sp = faiss.swig_ptr
- for i in range(ops.all_pts.size()):
- ps.update_bounds(cno, ops.all_pts.at(i), sp(bounds[1:2]), sp(bounds[0:1]))
- # lower_bound_t, upper_bound_perf
- return bounds[0], bounds[1]
-
-
-
-def explore_parameter_space_range(index, xq, gt, ps, radius):
- """ exploration of the parameter space for range search, using the
- Average Precision as criterion
- """
-
- n_experiments = ps.n_experiments
- n_comb = ps.n_combinations()
- min_time = ps.min_test_duration
- verbose = ps.verbose
-
- gt_nres, gt_I, gt_D = gt
- gt_lims = np.zeros(len(gt_nres) + 1, dtype=int)
- gt_lims[1:] = np.cumsum(gt_nres)
- gt = (gt_lims, gt_I, gt_D)
-
- ops = faiss.OperatingPoints()
-
- def run_1_experiment(cno):
- ps.set_index_parameters(index, cno)
-
- nrun = 0
- t0 = time.time()
- while True:
- lims, D, I = index.range_search(xq, radius)
- nrun += 1
- t1 = time.time()
- if t1 - t0 > min_time:
- break
-
- t_search = (t1 - t0) / nrun
- perf = eval_range_search.compute_AP(gt, (lims, I, D))
- keep = ops.add(perf, t_search, ps.combination_name(cno), cno)
-
- return len(D), perf, t_search, nrun, keep
-
- if n_experiments == 0:
- # means exhaustive run
- for cno in range(n_comb):
- nres, perf, t_search, nrun, keep = run_1_experiment(cno)
-
- if verbose:
- print(" %d/%d: %s nres=%d perf=%.3f t=%.3f s %s" % (
- cno, n_comb,
- ps.combination_name(cno),
- nres, perf, t_search, "*" if keep else ""))
- return ops
-
- n_experiments = min(n_experiments, n_comb)
-
- perm = np.zeros(n_experiments, int)
- # make sure the slowest and fastest experiment are run
- perm[0] = 0
- perm[1] = n_comb - 1
- rs = np.random.RandomState(1234)
- perm[2:] = 1 + rs.choice(n_comb - 2, n_experiments - 2, replace=False)
-
- for xp, cno in enumerate(perm):
- cno = int(cno)
- if verbose:
- print(" %d/%d: cno=%d %s " % (
- xp, n_experiments, cno, ps.combination_name(cno)),
- end="", flush=True)
-
- # check if we can skip this experiment
- lower_bound_t, upper_bound_perf = op_compute_bounds(ps, ops, cno)
-
- best_t = ops.t_for_perf(upper_bound_perf)
-
- if verbose:
- print("bounds [perf<=%.3f t>=%.3f] " % (
- upper_bound_perf, lower_bound_t),
- end="skip\n" if best_t <= lower_bound_t else " "
- )
- if best_t <= lower_bound_t:
- continue
-
- nres, perf, t_search, nrun, keep = run_1_experiment(cno)
-
- if verbose:
- print(" nres %d perf %.3f t %.3f (%d %s) %s" % (
- nres, perf, t_search, nrun,
- "runs" if nrun >= 2 else "run",
- "*" if keep else ""))
-
- return ops
-
-
-####################################################################
-# Driver functions
-####################################################################
-
-
-
-def run_experiments_searchparams(ds, index, args):
- """
- Evaluate a predefined set of runtime parameters
- """
- k = args.k
- xq = ds.get_queries()
-
- nq = len(xq)
-
- ps = faiss.ParameterSpace()
- ps.initialize(index)
-
- header, _ = result_header(ds, args)
-
- searchparams = args.searchparams
-
- print(f"Running evaluation on {len(searchparams)} searchparams")
- print(header)
- maxw = max(max(len(p) for p in searchparams), 40)
- for params in searchparams:
- ps.set_index_parameters(index, params)
-
- print(params.ljust(maxw), end=' ')
- sys.stdout.flush()
-
- if ds.search_type() == "knn":
- eval_setting_knn(
- index, xq, ds.get_groundtruth(k=args.k),
- k=args.k,
- inter=args.inter, min_time=args.min_test_duration,
- query_bs=args.query_bs
- )
- else:
- eval_setting_range(
- index, xq, ds.get_groundtruth(k=args.k),
- radius=args.radius,
- inter=args.inter, min_time=args.min_test_duration,
- query_bs=args.query_bs
- )
-
-
-def run_experiments_autotune(ds, index, args):
- """ Explore the space of parameters and keep Pareto-optimal ones. """
- k = args.k
-
- xq = ds.get_queries()
- nq = len(xq)
-
- ps = faiss.ParameterSpace()
- ps.initialize(index)
-
- ps.n_experiments = args.n_autotune
- ps.min_test_duration = args.min_test_duration
-
- for kv in args.autotune_max:
- k, vmax = kv.split(':')
- vmax = float(vmax)
- print("limiting %s to %g" % (k, vmax))
- pr = ps.add_range(k)
- values = faiss.vector_to_array(pr.values)
- values = np.array([v for v in values if v < vmax])
- faiss.copy_array_to_vector(values, pr.values)
-
- for kv in args.autotune_range:
- k, vals = kv.split(':')
- vals = np.fromstring(vals, sep=',')
- print("setting %s to %s" % (k, vals))
- pr = ps.add_range(k)
- faiss.copy_array_to_vector(vals, pr.values)
-
- header, crit = result_header(ds, args)
-
- # then we let Faiss find the optimal parameters by itself
- print("exploring operating points, %d threads" % faiss.omp_get_max_threads());
- ps.display()
-
- t0 = time.time()
-
- if ds.search_type() == "knn":
- # by default, the criterion will request only 1 NN
- crit.nnn = args.k
- gt_I, gt_D = ds.get_groundtruth(k=args.k)
- crit.set_groundtruth(None, gt_I.astype('int64'))
- op = ps.explore(index, xq, crit)
- elif ds.search_type() == "range":
- op = explore_parameter_space_range(
- index, xq, ds.get_groundtruth(), ps, args.radius
- )
- else:
- assert False
-
- print("Done in %.3f s, available OPs:" % (time.time() - t0))
- op.display()
-
- print("Re-running evaluation on selected OPs")
- print(header)
- opv = op.optimal_pts
- maxw = max(max(len(opv.at(i).key) for i in range(opv.size())), 40)
- for i in range(opv.size()):
- opt = opv.at(i)
-
- ps.set_index_parameters(index, opt.key)
-
- print(opt.key.ljust(maxw), end=' ')
- sys.stdout.flush()
- if ds.search_type() == "knn":
- eval_setting_knn(
- index, xq, ds.get_groundtruth(k=args.k),
- k=args.k,
- inter=args.inter, min_time=args.min_test_duration
- )
- else:
- eval_setting_range(
- index, xq, ds.get_groundtruth(k=args.k),
- radius=args.radius,
- inter=args.inter, min_time=args.min_test_duration
- )
-
-
-class DatasetWrapInPairwiseQuantization:
-
- def __init__(self, ds, C):
- self.ds = ds
- self.C = C
- self.Cq = np.linalg.inv(C.T)
- # xb_pw = np.ascontiguousarray((C @ xb.T).T)
- # xq_pw = np.ascontiguousarray((Cq @ xq.T).T)
- # copy fields
-
- for name in "nb d nq dtype distance search_type get_groundtruth".split():
- setattr(self, name, getattr(ds, name))
-
- def get_dataset(self):
- return self.ds.get_dataset() @ self.C.T
-
- def get_queries(self):
- return self.ds.get_queries() @ self.Cq.T
-
- def get_dataset_iterator(self, bs=512, split=(1,0)):
- for xb in self.ds.get_dataset_iterator(bs=bs, split=split):
- yield xb @ self.C.T
-
-
-####################################################################
-# Main
-####################################################################
-
-
-def main():
-
- parser = argparse.ArgumentParser()
-
- def aa(*args, **kwargs):
- group.add_argument(*args, **kwargs)
-
- group = parser.add_argument_group('What to do')
- aa('--build', default=False, action="store_true")
- aa('--search', default=False, action="store_true")
- aa('--prepare', default=False, action="store_true",
- help="call prepare() to download the dataset before computing")
-
- group = parser.add_argument_group('dataset options')
- aa('--dataset', choices=DATASETS.keys(), required=True)
- aa('--basedir', help="override basedir for dataset")
- aa('--pairwise_quantization', default="",
- help="load/store pairwise quantization matrix")
- aa('--query_bs', default=-1, type=int,
- help='perform queries in batches of this size')
-
- group = parser.add_argument_group('index construction')
-
- aa('--indexkey', default='HNSW32', help='index_factory type')
- aa('--by_residual', default=-1, type=int,
- help="set if index should use residuals (default=unchanged)")
- aa('--M0', default=-1, type=int, help='size of base level')
- aa('--maxtrain', default=0, type=int,
- help='maximum number of training points (0 to set automatically)')
- aa('--indexfile', default='', help='file to read or write index from')
- aa('--add_bs', default=100000, type=int,
- help='add elements index by batches of this size')
- aa('--add_splits', default=1, type=int,
- help="Do adds in this many splits (otherwise risk of OOM for large datasets)")
- aa('--stop_at_split', default=-1, type=int,
- help="stop at this split (for debugging)")
-
- aa('--no_precomputed_tables', action='store_true', default=False,
- help='disable precomputed tables (uses less memory)')
- aa('--clustering_niter', default=-1, type=int,
- help='number of clustering iterations (-1 = leave default)')
- aa('--two_level_clustering', action="store_true", default=False,
- help='perform a 2-level tree clustering')
- aa('--train_on_gpu', default=False, action='store_true',
- help='do training on GPU')
- aa('--quantizer_efConstruction', default=-1, type=int,
- help="override the efClustering of the quantizer")
- aa('--quantizer_add_efSearch', default=-1, type=int,
- help="override the efSearch of the quantizer at add time")
- aa('--buildthreads', default=-1, type=int,
- help='nb of threads to use at build time')
-
- group = parser.add_argument_group('searching')
-
- aa('--k', default=10, type=int, help='nb of nearest neighbors')
- aa('--radius', default=96237, type=float, help='radius for range search')
- aa('--inter', default=True, action='store_true',
- help='use intersection measure instead of 1-recall as metric')
- aa('--searchthreads', default=-1, type=int,
- help='nb of threads to use at search time')
- aa('--searchparams', nargs='+', default=['autotune'],
- help="search parameters to use (can be autotune or a list of params)")
- aa('--n_autotune', default=500, type=int,
- help="max nb of autotune experiments")
- aa('--autotune_max', default=[], nargs='*',
- help='set max value for autotune variables format "var:val" (exclusive)')
- aa('--autotune_range', default=[], nargs='*',
- help='set complete autotune range, format "var:val1,val2,..."')
- aa('--min_test_duration', default=3.0, type=float,
- help='run test at least for so long to avoid jitter')
- aa('--parallel_mode', default=-1, type=int,
- help="set search-time parallel mode for IVF indexes")
-
- group = parser.add_argument_group('computation options')
- aa("--maxRAM", default=-1, type=int, help="set max RSS in GB (avoid OOM crash)")
-
-
- args = parser.parse_args()
-
- print("args=", args)
-
- if args.basedir:
- print("setting datasets basedir to", args.basedir)
- benchmark.datasets.BASEDIR
- benchmark.datasets.BASEDIR = args.basedir
-
- if args.maxRAM > 0:
- print("setting max RSS to", args.maxRAM, "GiB")
- resource.setrlimit(
- resource.RLIMIT_DATA, (args.maxRAM * 1024 ** 3, resource.RLIM_INFINITY)
- )
-
- os.system('echo -n "nb processors "; '
- 'cat /proc/cpuinfo | grep ^processor | wc -l; '
- 'cat /proc/cpuinfo | grep ^"model name" | tail -1')
-
- ds = DATASETS[args.dataset]()
- print(ds)
-
- nq, d = ds.nq, ds.d
- nb, d = ds.nq, ds.d
-
- if args.prepare:
- print("downloading dataset...")
- ds.prepare()
- print("dataset ready")
-
- if not (args.build or args.search):
- return
-
- if args.pairwise_quantization:
- if os.path.exists(args.pairwise_quantization):
- print("loading pairwise quantization matrix", args.pairwise_quantization)
- C = np.load(args.pairwise_quantization)
- else:
- print("training pairwise quantization")
- xq_train = ds.get_query_train()
- G = xq_train.T @ xq_train
- C = np.linalg.cholesky(G).T
- print("store matrix in", args.pairwise_quantization)
- np.save(args.pairwise_quantization, C)
- # Cq = np.linalg.inv(C.T)
- # xb_pw = np.ascontiguousarray((C @ xb.T).T)
- # xq_pw = np.ascontiguousarray((Cq @ xq.T).T)
- ds = DatasetWrapInPairwiseQuantization(ds, C)
-
- if args.build:
- print("build index, key=", args.indexkey)
- index = build_index(args, ds)
- else:
- print("reading", args.indexfile)
- index = faiss.read_index(args.indexfile)
-
- index_ivf, vec_transform = unwind_index_ivf(index)
- if vec_transform is None:
- vec_transform = lambda x: x
-
- if index_ivf is not None:
- print("imbalance_factor=", index_ivf.invlists.imbalance_factor())
-
- if args.no_precomputed_tables:
- if isinstance(index_ivf, faiss.IndexIVFPQ):
- print("disabling precomputed table")
- index_ivf.use_precomputed_table = -1
- index_ivf.precomputed_table.clear()
-
- if args.indexfile:
- print("index size on disk: ", os.stat(args.indexfile).st_size)
-
- print("current RSS:", faiss.get_mem_usage_kb() * 1024)
-
- precomputed_table_size = 0
- if hasattr(index_ivf, 'precomputed_table'):
- precomputed_table_size = index_ivf.precomputed_table.size() * 4
-
- print("precomputed tables size:", precomputed_table_size)
-
- if args.search:
-
- if args.searchthreads == -1:
- print("Search threads:", faiss.omp_get_max_threads())
- else:
- print("Setting nb of threads to", args.searchthreads)
- faiss.omp_set_num_threads(args.searchthreads)
-
- if args.parallel_mode != -1:
- print("setting IVF parallel mode to", args.parallel_mode)
- index_ivf.parallel_mode
- index_ivf.parallel_mode = args.parallel_mode
-
- if args.searchparams == ["autotune"]:
- run_experiments_autotune(ds, index, args)
- else:
- run_experiments_searchparams(ds, index, args)
-
-
-if __name__ == "__main__":
- main()
diff --git a/track1_baseline_faiss/parse_results.py b/track1_baseline_faiss/parse_results.py
deleted file mode 100644
index 48e3411ac..000000000
--- a/track1_baseline_faiss/parse_results.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""
-Parse log files from baseline_faiss.py
-
-"""
-import os
-import numpy as np
-
-
-
-def parse_result_file(fname):
- # print fname
- st = 0
- res = []
- keys = []
- stats = {}
- stats['run_version'] = fname[-8]
- indexkey = None
- for l in open(fname):
- if l.startswith("srun:"):
- # looks like a crash...
- if indexkey is None:
- raise RuntimeError("instant crash")
- break
- elif st == 0:
- if l.startswith("dataset in dimension"):
- fi = l.split()
- stats["d"] = int(fi[3][:-1])
- stats["nq"] = int(fi[9])
- stats["nb"] = int(fi[11])
- stats["nt"] = int(fi[13])
- if l.startswith('index size on disk:'):
- stats['index_size'] = int(l.split()[-1])
- if l.startswith('current RSS:'):
- stats['RSS'] = int(l.split()[-1])
- if l.startswith('precomputed tables size:'):
- stats['tables_size'] = int(l.split()[-1])
- if l.startswith('Setting nb of threads to'):
- stats['n_threads'] = int(l.split()[-1])
- if l.startswith(' add in'):
- stats['add_time'] = float(l.split()[-2])
- if l.startswith('args:'):
- args = eval(l[l.find(' '):])
- indexkey = args.indexkey
- if l.startswith('build index, key='):
- indexkey = l.split()[-1]
- elif "time(ms/q)" in l:
- # result header
- if 'R@1 R@10 R@100' in l:
- stats["measure"] = "recall"
- stats["ranks"] = [1, 10, 100]
- elif 'I@1 I@10 I@100' in l:
- stats["measure"] = "inter"
- stats["ranks"] = [1, 10, 100]
- elif 'inter@' in l:
- stats["measure"] = "inter"
- fi = l.split()
- if fi[1] == "inter@":
- rank = int(fi[2])
- else:
- rank = int(fi[1][len("inter@"):])
- stats["ranks"] = [rank]
- elif 'AP' in l:
- stats["measure"] = "average_precision"
- else:
- assert False
- st = 1
- elif 'index size on disk:' in l:
- stats["index_size"] = int(l.split()[-1])
- elif st == 1:
- st = 2
- elif st == 2:
- fi = l.split()
- if l[0] == " ":
- # means there are 0 parameters
- fi = [""] + fi
- keys.append(fi[0])
- if len(fi[1:]) > 0:
- res.append([float(x) for x in fi[1:]])
- return indexkey, np.array(res), keys, stats
-
-
-def find_latest_version(fname):
- """ all log files are called
- XX.a.log
- XX.b.log
-
- Where XX is the experiment id and a, b... are versions.
- The version is used when the same experiment needs to be
- redone because it failed. This function returns the latest version
- """
- assert fname.endswith(".log")
- pref = fname[:-5]
- lv = ""
- for suf in "abcdefghijklmnopqrs":
- fname = pref + suf + '.log'
- if os.path.exists(fname):
- lv = fname
- assert lv
- return lv
\ No newline at end of file
diff --git a/track1_baseline_faiss/plots/bigann-1B.png b/track1_baseline_faiss/plots/bigann-1B.png
deleted file mode 100644
index 7a1d4a5b7..000000000
Binary files a/track1_baseline_faiss/plots/bigann-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/plots/deep-1B.png b/track1_baseline_faiss/plots/deep-1B.png
deleted file mode 100644
index 0c6bfc32e..000000000
Binary files a/track1_baseline_faiss/plots/deep-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/plots/msspacev-1B.png b/track1_baseline_faiss/plots/msspacev-1B.png
deleted file mode 100644
index a8f0c5fda..000000000
Binary files a/track1_baseline_faiss/plots/msspacev-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/plots/msturing-1B.png b/track1_baseline_faiss/plots/msturing-1B.png
deleted file mode 100644
index 2363edcd2..000000000
Binary files a/track1_baseline_faiss/plots/msturing-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/plots/ssnpp-1B.png b/track1_baseline_faiss/plots/ssnpp-1B.png
deleted file mode 100644
index 96e1b9720..000000000
Binary files a/track1_baseline_faiss/plots/ssnpp-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/plots/text2image-1B.png b/track1_baseline_faiss/plots/text2image-1B.png
deleted file mode 100644
index 5f7b9b302..000000000
Binary files a/track1_baseline_faiss/plots/text2image-1B.png and /dev/null differ
diff --git a/track1_baseline_faiss/run_baselines.bash b/track1_baseline_faiss/run_baselines.bash
deleted file mode 100644
index 62901423f..000000000
--- a/track1_baseline_faiss/run_baselines.bash
+++ /dev/null
@@ -1,530 +0,0 @@
-set -e
-
-export PYTHONPATH=.
-
-
-
-
-function run_on () {
- local sbatch_opt="$1"
- shift
- local name=$1
- shift
- local torun=" $@ "
-
- if [ -e slurm_scripts/$name.sh ]; then
- echo "script" slurm_scripts/$name.sh exists
- exit 1
- fi
-
- echo -n $name " "
-
- echo $@ > slurm_scripts/$name.sh
-
- sbatch $sbatch_opt \
- -J $name -o logs/$name.log \
- --wrap "bash slurm_scripts/$name.sh"
-
-}
-
-
-function run_on_1gpu () {
- run_on "--gres=gpu:1 --ntasks=1 --time=30:00:00 --cpus-per-task=20
- --partition=devlab --mem=64g --nodes=1 " "$@"
-}
-
-function run_on_1gpu_learnlab () {
- run_on "--gres=gpu:1 --ntasks=1 --time=30:00:00 --cpus-per-task=20
- --partition=learnlab --mem=64g --nodes=1 " "$@"
-}
-function run_on_half_machine () {
- run_on "--gres=gpu:1 --ntasks=1 --time=30:00:00 --cpus-per-task=40
- --partition=learnlab --mem=256g --nodes=1 " "$@"
-}
-
-function run_on_2gpu_ram256 () {
- run_on "--gres=gpu:2 --ntasks=1 --time=30:00:00 --cpus-per-task=20
- --partition=learnlab --mem=256g --nodes=1 " "$@"
-}
-
-
-
-##############################################################
-# Small scale experiments to evaluate effect of 2-level clustering
-##############################################################
-
-# compare 2-level 65k clustering index and regular one
-
-basedir=data/track1_baseline_faiss
-
-
-if false; then
-
-dsname=bigann-10M
-
-
-run_on_1gpu $dsname.IVF65k_HNSW.a \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.faissindex \
- --indexkey PCAR64,IVF65536_HNSW32,Flat --maxtrain $((65536 * 50)) \
- --search --train_on_gpu
-
-
-run_on_1gpu $dsname.IVF65k_2level_HNSW.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.IVF65k_2level_HNSW.faissindex \
- --indexkey PCAR64,IVF65536_HNSW32,Flat --maxtrain $((65536 * 50)) \
- --two_level_clustering \
- --search
-
-
-
-
-# for efC in 50 100 200; do
-
-for efC in 400 800; do
-
-run_on_1gpu $dsname.IVF65k_HNSW_efC$efC.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.IVF65k_HNSW_efC$efC.faissindex \
- --indexkey PCAR64,IVF65536_HNSW32,Flat --maxtrain $((65536 * 50)) \
- --quantizer_efConstruction $efC \
- --build --search --train_on_gpu
-
-done
-
-
-
-# for efS in 20 40 80; do
-for efS in 160 320; do
-
-name=$dsname.IVF65k_2level_HNSW_efC200_efS$efS
-
-run_on_1gpu $name.a \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey PCAR64,IVF65536_HNSW32,Flat --maxtrain $((65536 * 50)) \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch $efS \
- --two_level_clustering \
- --build --search
-
-done
-
-
-
-
-##############################################################
-# Experiments on scale 100M
-##############################################################
-
-# .a: build
-# .c: eval w 32 threads
-
-# start with least problematic datasets (no IP search, no range search)
-# msspace-1B may need to redo experiments because of ties in distance computations
-for dsname in bigann-100M deep-100M msturing-100M msspacev-100M; do
-
- for nc in 256k 1M; do
-
- case $nc in
- 1M) ncn=$((1<<20)) ;;
- 256k) ncn=$((1<<18)) ;;
- esac
-
- name=$dsname.IVF${nc}_2level_PQ32
-
- run_on_half_machine $name.c \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ32_128,IVF${ncn}_HNSW32,PQ32 \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --search --searchthreads 32 \
- --maxRAM 256
-
- name=$dsname.IVF${nc}_2level_PQ64x4fsr
-
- run_on_half_machine $name.c \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ64_128,IVF${ncn}_HNSW32,PQ64x4fsr \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --search --searchthreads 32 \
- --maxRAM 256
-
- done
-
-done
-
-
-##############################################################
-# Experiments on scale 1B
-##############################################################
-
-# .a: build
-# .b: eval w 32 threads
-# .c: redo bigann eval
-# .d: with ssnpp, forgot to build...
-# .f: redo t2i 64x4 (eval only)
-
-# start with least problematic datasets (no IP search, no range search)
-# msspace-1B may need to redo experiments because of ties in distance computations
-
-# for dsname in bigann-1B deep-1B msturing-1B msspacev-1B; do
-# for dsname in bigann-1B; do
-# for dsname in ssnpp-1B; do
-# for nc in 1M 4M; do
-
-fi
-
-for dsname in text2image-1B; do
-
- for nc in 1M; do
-
- case $nc in
- 1M) ncn=$((1<<20)) ;;
- 4M) ncn=$((1<<22)) ;;
- esac
-
- if false ;then
-
- name=$dsname.IVF${nc}_2level_PQ32
-
- run_on_half_machine $name.e \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ32_128,IVF${ncn}_HNSW32,PQ32 \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --build --search --searchthreads 32 \
- --maxRAM 256
-
- fi
- name=$dsname.IVF${nc}_2level_PQ64x4fsr
-
- run_on_half_machine $name.g \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ64_128,IVF${ncn}_HNSW32,PQ64x4fsr \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --search --searchthreads 32 \
- --maxRAM 256 --autotune_max nprobe:513
-
-
- done
-
-done
-
-if false; then
-
-# speed up construction
-
-dsname=ssnpp-1B
-nc=1M
-ncn=$((1<<20))
-
-name=$dsname.IVF${nc}_2level_aefS40_PQ32
-
-un_on_half_machine $name.d \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ32_128,IVF${ncn}_HNSW32,PQ32 \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 40 \
- --add_splits 30 \
- --two_level_clustering \
- --build --search --searchthreads 32 \
- --maxRAM 256
-
-
-# find a way to not OOM during autotune
-
-function ssnpp_no_OOM () {
- local key=$1
- shift
- dsname=ssnpp-1B
- nc=1M
- ncn=$((1<<20))
-
- name=$dsname.IVF${nc}_2level_PQ32.search.$key
-
- run_on_half_machine $name.a \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.IVF${nc}_2level_PQ32.faissindex \
- --search --searchthreads 32 \
- --maxRAM 256 "$@"
-}
-
-ssnpp_no_OOM radius75000 --radius 75000
-ssnpp_no_OOM radius80000 --radius 80000
-ssnpp_no_OOM radius60000 --radius 60000
-ssnpp_no_OOM maxNP1024 --autotune_max nprobe:1025
-ssnpp_no_OOM maxEFS256 --autotune_max quantizer_efSearch:257
-
-
-
-##############################################################
-# Experiments with 64 bytes per vector
-##############################################################
-
-# .a: initial run and build
-# .b: re-run to get more detailed search stats
-
-
-for dsname in bigann-1B deep-1B msturing-1B msspacev-1B; do
- nc=1M
- ncn=$((1<<20))
-
- name=$dsname.IVF${nc}_2level_PQ64
-
- run_on_half_machine $name.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ64_128,IVF${ncn}_HNSW32,PQ64 \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --search --searchthreads 32 \
- --maxRAM 256
-
- name=$dsname.IVF${nc}_2level_PQ128x4fsr
-
- run_on_half_machine $name.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$name.faissindex \
- --indexkey OPQ128_128,IVF${ncn}_HNSW32,PQ128x4fsr \
- --maxtrain 100000000 \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --two_level_clustering \
- --search --searchthreads 32 \
- --maxRAM 256
-
-done
-
-
-
-##############################################################
-# 10M scale exeperiment for text2image
-##############################################################
-
-dsname=text2image-10M
-
-
-for nc in 16k 65k; do
-
- case $nc in
- 16k) ncn=$((1<<14)) ;;
- 65k) ncn=$((1<<16)) ;;
- esac
-
- # baseline
- key=IVF$nc
- run_on_1gpu $dsname.$key.d \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey IVF${ncn},Flat --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
- # loss due to 2-level
- key=IVF${nc}_2level
- run_on_1gpu $dsname.$key.d \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey IVF${ncn},Flat --maxtrain $((ncn * 4 * 50)) \
- --build --search --two_level_clustering
-
- # loss due to HNSW
- key=IVF${nc}_HNSW
- run_on_1gpu $dsname.$key.d \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey IVF${ncn}_HNSW32,Flat --maxtrain $((ncn * 4 * 50)) \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --build --search --train_on_gpu
-
- # loss due to 2-level + HNSW
- key=IVF${nc}_2level_HNSW
- run_on_1gpu $dsname.$key.d \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey IVF${ncn}_HNSW32,Flat --maxtrain $((ncn * 4 * 50)) \
- --quantizer_efConstruction 200 \
- --quantizer_add_efSearch 80 \
- --build --search --two_level_clustering
-
-done
-
-# evaluate various IVF codes
-
-ncn=16384
-
-
-key=IVF16k,SQ8
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR200,IVF16384,SQ8 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,SQ8_nores
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR200,IVF16384,SQ8 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu --by_residual 0
-
-key=IVF16k,SQ6
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR200,IVF16384,SQ6 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,SQ6_nores
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR200,IVF16384,SQ6 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu --by_residual 0
-
-
-key=IVF16k,SQ8_PQ32
-run_on_1gpu_learnlab $dsname.$key.a \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey OPQ32_128,IVF16384,PQ32 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,SQ8_PQ32_nores
-run_on_1gpu_learnlab $dsname.$key.a \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey OPQ32_128,IVF16384,PQ32 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu --by_residual 0
-
-
-key=IVF16k,SQ4
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR200,IVF16384,SQ4 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,SQ4_PCAR100
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey PCAR100,IVF16384,SQ4 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,RR192_PQ32
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR192,IVF16384,PQ32 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-key=IVF16k,RR192_PQ32x12
-run_on_1gpu_learnlab $dsname.$key.b \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey RR192,IVF16384,PQ32x12 --maxtrain $((ncn * 4 * 50)) \
- --build --search --train_on_gpu
-
-
-dsname=text2image-10M
-
-key=IVF16k,PQ48
-run_on_1gpu $dsname.$key.c \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey OPQ48_192,IVF16384,PQ48 --maxtrain $((65536 * 50)) \
- --search --train_on_gpu
-
-key=IVF16k,PQ64
-run_on_1gpu $dsname.$key.c \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey OPQ64_192,IVF16384,PQ64 --maxtrain $((65536 * 50)) \
- --search --train_on_gpu
-
-
-dsname=text2image-10M
-key=IVF16k,PQ48
-run_on_1gpu $dsname.$key.c \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey OPQ48_192,IVF16384,PQ48 --maxtrain $((65536 * 50)) \
- --search --train_on_gpu
-
-
-
-## try out additive quants
-
-
-export PYTHONPATH=/private/home/matthijs/faiss_versions/faiss_add_quant_search/build/faiss/python/build/lib:.
-dsname=text2image-10M
-for key in IVF16384,RQ32x8_Nfloat IVF16384,RQ31x8_Nqint8 IVF16384,LSQ32x8_Nfloat IVF16384,LSQ31x8_Nqint8 ; do
-run_on_1gpu $dsname.$key.g \
- python -u track1_baseline_faiss/baseline_faiss.py \
- --dataset $dsname --indexfile $basedir/$dsname.$key.faissindex \
- --indexkey $key --maxtrain $((65536 * 50)) \
- --search --build # --train_on_gpu
-
-done
-
-
-##############################################################
-# GPU based search (T3)
-##############################################################
-
-basedir=data/track3_baseline_faiss
-dsname=deep-1B
-
-#.a: run with a too tight limit in RAM
-#.b: increased RAM
-
-key=IVF262k,PQ8
-run_on_2gpu_ram256 T3.$dsname.$key.b \
- python -u track3_baseline_faiss/gpu_baseline_faiss.py \
- --maxRAM 256 \
- --dataset $dsname --indexkey IVF$((1<<18)),SQ8 \
- --build \
- --searchparams nprobe={1,4,16,64,256,1024} \
- --train_on_gpu --quantizer_on_gpu_add \
- --indexfile $basedir/$dsname.$key.faissindex \
- --add_splits 30 \
- --search \
- --parallel_mode 3 --quantizer_on_gpu_search
-
-
-key=IVF1M,PQ8
-run_on_2gpu_ram256 T3.$dsname.$key.b \
- python -u track3_baseline_faiss/gpu_baseline_faiss.py \
- --maxRAM 256 \
- --dataset $dsname --indexkey IVF$((1<<20)),SQ8 \
- --build \
- --searchparams nprobe={1,4,16,64,256,1024} \
- --train_on_gpu --quantizer_on_gpu_add \
- --indexfile $basedir/$dsname.$key.faissindex \
- --add_splits 30 \
- --search \
- --parallel_mode 3 --quantizer_on_gpu_search
-
-
-fi
\ No newline at end of file
diff --git a/track3_baseline_faiss/README.md b/track3_baseline_faiss/README.md
deleted file mode 100644
index 903e6b9fb..000000000
--- a/track3_baseline_faiss/README.md
+++ /dev/null
@@ -1,73 +0,0 @@
-# Running the Faiss GPU baseline
-
-The script here is based on the T1 baseline, so please take a look at [the Track 1 baseline](../track1_baseline_faiss/README.md) first.
-
-## Installing software
-
-See [this doc](../track1_baseline_faiss/README.md#installing-software) but instead of installing faiss-cpu, use:
-
-```
-conda install -c pytorch faiss-gpu cudatoolkit=10.2
-```
-
-## How to use the GPU
-
-This script focuses on exploiting the GPU for coarse quantization.
-Therefore, it is suitable for large codebooks.
-
-The GPU can be used in the following phases:
-
-- training: `--train_with_gpu` will move the training of the coarse quantizer to GPU
-
-- vector adding to the index: `--quantizer_on_gpu_add --` will do the adding on GPU
-
-- search: `--parallel_mode 3 --quantizer_on_gpu_search` will do coarse quantization on GPU at search time.
-
-## Building the index and searching
-
-The hardware environment is: 1 GPU on a machine with 768G RAM (practically unlimited).
-Therefore, the approach is to do the coarse quantization on GPU and store the IVF index in RAM with mild compression (PQ8).
-This means that to keep the GPU busy the number of centroids should be as large as possible.
-We use 1M in the example below.
-The GPU uses brute force computations to find the nearest centroids.
-
-### 100M-scale
-
-The following command runs the index constuction and evaluates the search performance:
-
-```bash
-python track3_baseline_faiss/gpu_baseline_faiss.py --dataset deep-100M \
- --indexkey IVF65536,SQ8 \
- --train_on_gpu \
- --build --quantizer_on_gpu_add --add_splits 30 \
- --search \
- --searchparams nprobe={1,4,16,64,256} \
- --parallel_mode 3 --quantizer_on_gpu_search
-```
-
-Example logs [without GPU](https://gist.github.com/mdouze/9e000be47c499f79aaec0166365ef654) and [with GPU](https://gist.github.com/mdouze/cd14c802b924299aa2a92db6e05df857) at search time.
-
-
-### 1B-scale
-
-```bash
-python track3_baseline_faiss/gpu_baseline_faiss.py --dataset deep-1B \
- --indexkey IVF$((1<<18)),SQ8 \
- --train_on_gpu \
- --build --quantizer_on_gpu_add --add_splits 30 \
- --search \
- --searchparams nprobe={1,4,16,64,256} \
- --parallel_mode 3 --quantizer_on_gpu_search
-```
-
-
-For the SSNPP dataset, please use `--parallel_mode 2` instead.
-
-
-### Results
-
-Similar to the track 1 results, we can plot the GPU search results in a plot of recall@10 vs. QPS.
-
-
-
-Caveat: here the GPU uses 20 CPU threads vs. 32 for the CPU, and the search is actually performed on 2 GPUs.
diff --git a/track3_baseline_faiss/gpu_baseline_faiss.py b/track3_baseline_faiss/gpu_baseline_faiss.py
deleted file mode 100644
index a3fc2d1e8..000000000
--- a/track3_baseline_faiss/gpu_baseline_faiss.py
+++ /dev/null
@@ -1,574 +0,0 @@
-import os
-import sys
-import time
-import pdb
-import gc
-import numpy as np
-import faiss
-import argparse
-import resource
-import threading
-from multiprocessing.pool import ThreadPool
-
-import benchmark.datasets
-from benchmark.datasets import DATASETS
-from benchmark.plotting import eval_range_search
-
-
-
-def unwind_index_ivf(index):
- if isinstance(index, faiss.IndexPreTransform):
- assert index.chain.size() == 1
- vt = index.chain.at(0)
- index_ivf, vt2 = unwind_index_ivf(faiss.downcast_index(index.index))
- assert vt2 is None
- return index_ivf, vt
- if hasattr(faiss, "IndexRefine") and isinstance(index, faiss.IndexRefine):
- return unwind_index_ivf(faiss.downcast_index(index.base_index))
- if isinstance(index, faiss.IndexIVF):
- return index, None
- else:
- return None, None
-
-def rate_limited_iter(l):
- 'a thread pre-processes the next element'
- pool = ThreadPool(1)
- res = None
-
- def next_or_None():
- try:
- return next(l)
- except StopIteration:
- return None
-
- while True:
- res_next = pool.apply_async(next_or_None)
- if res is not None:
- res = res.get()
- if res is None:
- return
- yield res
- res = res_next
-
-
-def build_index(args, ds):
- nq, d = ds.nq, ds.d
- nb, d = ds.nq, ds.d
-
- if args.buildthreads == -1:
- print("Build-time number of threads:", faiss.omp_get_max_threads())
- else:
- print("Set build-time number of threads:", args.buildthreads)
- faiss.omp_set_num_threads(args.buildthreads)
-
- metric_type = (
- faiss.METRIC_L2 if ds.distance() == "euclidean" else
- faiss.METRIC_INNER_PRODUCT if ds.distance() in ("ip", "angular") else
- 1/0
- )
- index = faiss.index_factory(d, args.indexkey, metric_type)
-
- index_ivf, vec_transform = unwind_index_ivf(index)
- if vec_transform is None:
- vec_transform = lambda x: x
- else:
- vec_transform = faiss.downcast_VectorTransform(vec_transform)
-
- if args.by_residual != -1:
- by_residual = args.by_residual == 1
- print("setting by_residual = ", by_residual)
- index_ivf.by_residual # check if field exists
- index_ivf.by_residual = by_residual
-
- if index_ivf:
- print("Update add-time parameters")
- # adjust default parameters used at add time for quantizers
- # because otherwise the assignment is inaccurate
- quantizer = faiss.downcast_index(index_ivf.quantizer)
- if isinstance(quantizer, faiss.IndexRefine):
- print(" update quantizer k_factor=", quantizer.k_factor, end=" -> ")
- quantizer.k_factor = 32 if index_ivf.nlist < 1e6 else 64
- print(quantizer.k_factor)
- base_index = faiss.downcast_index(quantizer.base_index)
- if isinstance(base_index, faiss.IndexIVF):
- print(" update quantizer nprobe=", base_index.nprobe, end=" -> ")
- base_index.nprobe = (
- 16 if base_index.nlist < 1e5 else
- 32 if base_index.nlist < 4e6 else
- 64)
- print(base_index.nprobe)
-
- index.verbose = True
- if index_ivf:
- index_ivf.verbose = True
- index_ivf.quantizer.verbose = True
- index_ivf.cp.verbose = True
-
-
- maxtrain = args.maxtrain
- if maxtrain == 0:
- if 'IMI' in args.indexkey:
- maxtrain = int(256 * 2 ** (np.log2(index_ivf.nlist) / 2))
- elif index_ivf:
- maxtrain = 50 * index_ivf.nlist
- else:
- # just guess...
- maxtrain = 256 * 100
- maxtrain = max(maxtrain, 256 * 100)
- print("setting maxtrain to %d" % maxtrain)
-
- # train on dataset
- print(f"getting first {maxtrain} dataset vectors for training")
-
- xt2 = next(ds.get_dataset_iterator(bs=maxtrain))
-
- print("train, size", xt2.shape)
- assert np.all(np.isfinite(xt2))
-
- t0 = time.time()
-
- if (isinstance(vec_transform, faiss.OPQMatrix) and
- isinstance(index_ivf, faiss.IndexIVFPQFastScan)):
- print(" Forcing OPQ training PQ to PQ4")
- ref_pq = index_ivf.pq
- training_pq = faiss.ProductQuantizer(
- ref_pq.d, ref_pq.M, ref_pq.nbits
- )
- vec_transform.pq
- vec_transform.pq = training_pq
-
- if args.clustering_niter >= 0:
- print(("setting nb of clustering iterations to %d" %
- args.clustering_niter))
- index_ivf.cp.niter = args.clustering_niter
-
- train_index = None
- if args.train_on_gpu:
- print("add a training index on GPU")
- train_index = faiss.index_cpu_to_all_gpus(
- faiss.IndexFlatL2(index_ivf.d))
- index_ivf.clustering_index = train_index
-
- index.train(xt2)
- print(" Total train time %.3f s" % (time.time() - t0))
-
- if train_index is not None:
- del train_index
- index_ivf.clustering_index = None
- gc.collect()
-
- print("adding")
-
- t0 = time.time()
-
- if not args.quantizer_on_gpu_add:
- i0 = 0
- for xblock in ds.get_dataset_iterator(bs=args.add_bs):
- i1 = i0 + len(xblock)
- print(" adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
- i0, i1, ds.nb, time.time() - t0,
- faiss.get_mem_usage_kb()))
- index.add(xblock)
- i0 = i1
- elif True:
- quantizer_gpu = faiss.index_cpu_to_all_gpus(index_ivf.quantizer)
-
- nsplit = args.add_splits
-
- def produce_batches(sno):
- for xblock in ds.get_dataset_iterator(bs=args.add_bs, split=(nsplit, sno)):
- _, assign = quantizer_gpu.search(xblock, 1)
- yield xblock, assign.ravel()
-
- i0 = 0
- for sno in range(nsplit):
- print(f"============== SPLIT {sno}/{nsplit}")
-
- stage2 = rate_limited_iter(produce_batches(sno))
- for xblock, assign in stage2:
- i1 = i0 + len(xblock)
- print(" adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
- i0, i1, ds.nb, time.time() - t0,
- faiss.get_mem_usage_kb()))
- index.add_core(
- len(xblock),
- faiss.swig_ptr(xblock),
- None,
- faiss.swig_ptr(assign)
- )
- i0 = i1
- del quantizer_gpu
- gc.collect()
-
-
- print(" add in %.3f s" % (time.time() - t0))
- if args.indexfile:
- print("storing", args.indexfile)
- faiss.write_index(index, args.indexfile)
-
- return index
-
-
-def compute_inter(a, b):
- nq, rank = a.shape
- ninter = sum(
- np.intersect1d(a[i, :rank], b[i, :rank]).size
- for i in range(nq)
- )
- return ninter / a.size
-
-
-
-def eval_setting_knn(index, xq, gt, k, inter, min_time):
- nq = xq.shape[0]
- gt_I, gt_D = gt
- ivf_stats = faiss.cvar.indexIVF_stats
- ivf_stats.reset()
- nrun = 0
-
- t0 = time.time()
- while True:
- D, I = index.search(xq, k)
- nrun += 1
- t1 = time.time()
- if t1 - t0 > min_time:
- break
- ms_per_query = ((t1 - t0) * 1000.0 / nq / nrun)
- if inter:
- rank = k
- inter_measure = compute_inter(gt[:, :rank], I[:, :rank])
- print("%.4f" % inter_measure, end=' ')
- else:
- for rank in 1, 10, 100:
- n_ok = (I[:, :rank] == gt[:, :1]).sum()
- print("%.4f" % (n_ok / float(nq)), end=' ')
- print(" %9.5f " % ms_per_query, end=' ')
-
- if ivf_stats.search_time == 0:
- # happens for IVFPQFastScan where the stats are not logged by default
- print("%12d %5.2f " % (ivf_stats.ndis / nrun, 0.0), end=' ')
- else:
- pc_quantizer = ivf_stats.quantization_time / ivf_stats.search_time * 100
- print("%12d %5.2f " % (ivf_stats.ndis / nrun, pc_quantizer), end=' ')
- print(nrun)
-
-
-def eval_setting_range(index, xq, gt, radius=0, inter=False, min_time=3.0, query_bs=-1):
- nq = xq.shape[0]
- gt_nres, gt_I, gt_D = gt
- gt_lims = np.zeros(nq + 1, dtype=int)
- gt_lims[1:] = np.cumsum(gt_nres)
- ivf_stats = faiss.cvar.indexIVF_stats
- ivf_stats.reset()
- nrun = 0
- t0 = time.time()
- while True:
- lims, D, I = index.range_search(xq, radius)
- nrun += 1
- t1 = time.time()
- if t1 - t0 > min_time:
- break
- ms_per_query = ((t1 - t0) * 1000.0 / nq / nrun)
-
- ap = eval_range_search.compute_AP((gt_lims, gt_I, gt_D), (lims, I, D))
- print("%.4f" % ap, end=' ')
- print(" %9.5f " % ms_per_query, end=' ')
-
- print("%12d %5d " % (ivf_stats.ndis / nrun, D.size), end=' ')
- print(nrun)
-
-
-class IndexQuantizerOnGPU:
- """ run query quantization on GPU """
-
- def __init__(self, index, search_bs):
- self.search_bs = search_bs
- index_ivf, vec_transform = unwind_index_ivf(index)
- self.index_ivf = index_ivf
- self.vec_transform = vec_transform
- self.quantizer_gpu = faiss.index_cpu_to_all_gpus(self.index_ivf.quantizer)
-
-
- def produce_batches(self, x, bs):
- n = len(x)
- nprobe = self.index_ivf.nprobe
- ivf_stats = faiss.cvar.indexIVF_stats
- for i0 in range(0, n, bs):
- xblock = x[i0:i0 + bs]
- t0 = time.time()
- D, I = self.quantizer_gpu.search(xblock, nprobe)
- ivf_stats.quantization_time += 1000 * (time.time() - t0)
- yield i0, xblock, D, I
-
-
- def search(self, x, k):
- bs = self.search_bs
- if self.vec_transform:
- x = self.vec_transform(x)
- nprobe = self.index_ivf.nprobe
- n, d = x.shape
- assert self.index_ivf.d == d
- D = np.empty((n, k), dtype=np.float32)
- I = np.empty((n, k), dtype=np.int64)
-
- sp = faiss.swig_ptr
- stage2 = rate_limited_iter(self.produce_batches(x, bs))
- t0 = time.time()
- for i0, xblock, Dc, Ic in stage2:
- ni = len(xblock)
- self.index_ivf.search_preassigned(
- ni, faiss.swig_ptr(xblock),
- k, sp(Ic), sp(Dc),
- sp(D[i0:]), sp(I[i0:]),
- False
- )
-
- return D, I
-
- def range_search(self, x, radius):
- bs = self.search_bs
- if self.vec_transform:
- x = self.vec_transform(x)
- nprobe = self.index_ivf.nprobe
- n, d = x.shape
- assert self.index_ivf.d == d
-
- sp = faiss.swig_ptr
- rsp = faiss.rev_swig_ptr
- stage2 = rate_limited_iter(self.produce_batches(x, bs))
- t0 = time.time()
- all_res = []
- nres = 0
- for i0, xblock, Dc, Ic in stage2:
- ni = len(xblock)
- res = faiss.RangeSearchResult(ni)
-
- self.index_ivf.range_search_preassigned(
- ni, faiss.swig_ptr(xblock),
- radius, sp(Ic), sp(Dc),
- res
- )
- all_res.append((ni, res))
- lims = rsp(res.lims, ni + 1)
- nres += lims[-1]
- nres = int(nres)
- lims = np.zeros(n + 1, int)
- I = np.empty(nres, int)
- D = np.empty(nres, 'float32')
-
- n0 = 0
- for ni, res in all_res:
- lims_i = rsp(res.lims, ni + 1)
- nd = int(lims_i[-1])
- Di = rsp(res.distances, nd)
- Ii = rsp(res.labels, nd)
- i0 = int(lims[n0])
- lims[n0: n0 + ni + 1] = lims_i + i0
- I[i0:i0 + nd] = Ii
- D[i0:i0 + nd] = Di
- n0 += ni
-
- return lims, D, I
-
-
-def run_experiments_searchparams(ds, index, args):
- k = args.k
-
- xq = ds.get_queries()
-
- nq = len(xq)
-
- ps = faiss.ParameterSpace()
- ps.initialize(index)
-
-
- # setup the Criterion object
- if args.inter:
- print("Optimize for intersection @ ", args.k)
- header = (
- '%-40s inter@%3d time(ms/q) nb distances %%quantization #runs' %
- ("parameters", args.k)
- )
- else:
- print("Optimize for 1-recall @ 1")
- header = (
- '%-40s R@1 R@10 R@100 time(ms/q) nb distances %%quantization #runs' %
- "parameters"
- )
-
- searchparams = args.searchparams
-
- print(f"Running evaluation on {len(searchparams)} searchparams")
- print(header)
- maxw = max(max(len(p) for p in searchparams), 40)
-
- if args.quantizer_on_gpu_search:
- index_wrap = IndexQuantizerOnGPU(index, args.search_bs)
- else:
- index_wrap = index
-
- for params in searchparams:
- ps.set_index_parameters(index, params)
-
- print(params.ljust(maxw), end=' ')
- sys.stdout.flush()
-
- if ds.search_type() == "knn":
- eval_setting_knn(
- index_wrap, xq, ds.get_groundtruth(k=args.k),
- k=args.k, inter=args.inter, min_time=args.min_test_duration
- )
- else:
- eval_setting_range(
- index_wrap, xq, ds.get_groundtruth(),
- radius=args.radius, inter=args.inter,
- min_time=args.min_test_duration
- )
-
-
-
-def main():
-
- parser = argparse.ArgumentParser()
-
- def aa(*args, **kwargs):
- group.add_argument(*args, **kwargs)
-
- group = parser.add_argument_group('What to do')
- aa('--build', default=False, action="store_true")
- aa('--search', default=False, action="store_true")
- aa('--prepare', default=False, action="store_true",
- help="call prepare() to download the dataset before computing")
-
- group = parser.add_argument_group('dataset options')
- aa('--dataset', choices=DATASETS.keys(), required=True)
- aa('--basedir', help="override basedir for dataset")
-
- group = parser.add_argument_group('index consturction')
-
- aa('--indexkey', default='IVF1204,Flat', help='index_factory type')
- aa('--by_residual', default=-1, type=int,
- help="set if index should use residuals (default=unchanged)")
- aa('--maxtrain', default=0, type=int,
- help='maximum number of training points (0 to set automatically)')
- aa('--indexfile', default='', help='file to read or write index from')
- aa('--add_bs', default=100000, type=int,
- help='add elements index by batches of this size')
- aa('--no_precomputed_tables', action='store_true', default=False,
- help='disable precomputed tables (uses less memory)')
- aa('--clustering_niter', default=-1, type=int,
- help='number of clustering iterations (-1 = leave default)')
- aa('--train_on_gpu', default=False, action='store_true',
- help='do training on GPU')
- aa('--buildthreads', default=-1, type=int,
- help='nb of threads to use at build time')
- aa('--quantizer_on_gpu_add', action="store_true", default=False,
- help="use GPU coarse quantizer at add time")
- aa('--add_splits', default=1, type=int,
- help="Do adds in this many splits (otherwise risk of OOM with GPU based adds)")
-
- group = parser.add_argument_group('searching')
-
- aa('--k', default=10, type=int, help='nb of nearest neighbors')
- aa('--radius', default=96237, type=float, help='radius for range search')
- aa('--inter', default=True, action='store_true',
- help='use intersection measure instead of 1-recall as metric')
- aa('--searchthreads', default=-1, type=int,
- help='nb of threads to use at search time')
- aa('--searchparams', nargs='+', default=['autotune'],
- help="search parameters to use (can be autotune or a list of params)")
- aa('--min_test_duration', default=3.0, type=float,
- help='run test at least for so long to avoid jitter')
- aa('--quantizer_on_gpu_search', action="store_true", default=False,
- help="use GPU coarse quantizer at search time")
- aa('--parallel_mode', default=-1, type=int,
- help="set search-time parallel mode for IVF indexes")
- aa('--search_bs', default=8192, type=int,
- help='search time batch size (for GPU/CPU tiling)')
-
- group = parser.add_argument_group('computation options')
- aa("--maxRAM", default=-1, type=int, help="set max RSS in GB (avoid OOM crash)")
-
-
- args = parser.parse_args()
-
- if args.basedir:
- print("setting datasets basedir to", args.basedir)
- benchmark.datasets.BASEDIR
- benchmark.datasets.BASEDIR = args.basedir
-
- if args.maxRAM > 0:
- print("setting max RSS to", args.maxRAM, "GiB")
- resource.setrlimit(
- resource.RLIMIT_DATA, (args.maxRAM * 1024 ** 3, resource.RLIM_INFINITY)
- )
-
- os.system('echo -n "nb processors "; '
- 'cat /proc/cpuinfo | grep ^processor | wc -l; '
- 'cat /proc/cpuinfo | grep ^"model name" | tail -1')
-
- ds = DATASETS[args.dataset]()
- print(ds)
-
- nq, d = ds.nq, ds.d
- nb, d = ds.nq, ds.d
-
- if args.prepare:
- print("downloading dataset...")
- ds.prepare()
- print("dataset ready")
-
- if not (args.build or args.search):
- return
-
- if args.build:
- print("build index, key=", args.indexkey)
- index = build_index(args, ds)
- else:
- print("reading", args.indexfile)
- index = faiss.read_index(args.indexfile)
-
- index_ivf, vec_transform = unwind_index_ivf(index)
- if vec_transform is None:
- vec_transform = lambda x: x
-
- if index_ivf is not None:
- print("imbalance_factor=", index_ivf.invlists.imbalance_factor())
-
- if args.no_precomputed_tables:
- if isinstance(index_ivf, faiss.IndexIVFPQ):
- print("disabling precomputed table")
- index_ivf.use_precomputed_table = -1
- index_ivf.precomputed_table.clear()
-
- if args.indexfile:
- print("index size on disk: ", os.stat(args.indexfile).st_size)
-
- print("current RSS:", faiss.get_mem_usage_kb() * 1024)
-
- precomputed_table_size = 0
- if hasattr(index_ivf, 'precomputed_table'):
- precomputed_table_size = index_ivf.precomputed_table.size() * 4
-
- print("precomputed tables size:", precomputed_table_size)
-
- if args.search:
-
- if args.searchthreads == -1:
- print("Search threads:", faiss.omp_get_max_threads())
- else:
- print("Setting nb of threads to", args.searchthreads)
- faiss.omp_set_num_threads(args.searchthreads)
-
- if args.parallel_mode != -1:
- print("setting IVF parallel mode to", args.parallel_mode)
- index_ivf.parallel_mode
- index_ivf.parallel_mode = args.parallel_mode
-
- if args.searchparams == ["autotune"]:
- run_experiments_autotune(ds, index, args)
- else:
- run_experiments_searchparams(ds, index, args)
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/track3_baseline_faiss/plots/T3_deep-1B.png b/track3_baseline_faiss/plots/T3_deep-1B.png
deleted file mode 100644
index d6902d6a0..000000000
Binary files a/track3_baseline_faiss/plots/T3_deep-1B.png and /dev/null differ
+ 
+ 
+ 
+ 
+ Harsha Simhadri
+ Microsoft Research India
+
+
+
+
+ 
+ Martin Aumüller
+ IT University of Copenhagen
+
+
+
+
+ 
+ Dmitry Baranchuk
+ Yandex
+
+
+
+
+ 
+ Matthijs Douze
+ Meta AI Research
+
+
+
+
+ 
+ Amir Ingber
+ Pinecone
+
+
+
+
+ 
+ Edo Liberty
+ Pinecone
+
+
+
+
+ 
+ Frank Liu
+ Zilliz
+
+
+
+
+ 
+ George Williams
+ Smile Identity
+
+