annbench: a lightweight benchmark for approximate nearest neighbor search

annbench is a simple benchmark for approximate nearest neighbor search algorithms in Python. This repository design is strongly influenced by a great project, ann-benchmarks, that provides comprehensive and thorough benchmarks for various algorithms. In contrast, we aim to deliver more lightweight and straightforward scripts with the following features.

Support Euclidean distance only
Support Recall@1 only
Support libraries installable via pip/conda only
Search with a single thread
Sweep by a single query parameter

Leaderboard

AWS Deep Learning AMI, c5.4xlarge (Aug 28, 2021)

Getting started

git clone https://github.com/matsui528/annbench.git
cd annbench
pip install -r requirements.txt
# conda install faiss-cpu -y -c pytorch  # If you'd like to try faiss, run this on anaconda
# conda install faiss-gpu -y -c pytorch  # or, if you have GPUs, install faiss-gpu

python download.py dataset=siftsmall  # Downloaded on ./dataset

python run.py dataset=siftsmall algo=annoy  # Indices are on ./interim. Results are on ./output

python plot.py   # Plots are on ./result_img

Run all algorithms on all dataset

# Downloading deep1m takes some hours
python download.py --multirun dataset=siftsmall,sift1m,deep1m

# Will take some hours
python run.py --multirun dataset=siftsmall,sift1m,deep1m algo=linear,annoy,ivfpq,hnsw,ivfpq4bit,scann,pq,pq4bit,hnsw_faiss,nsg
# Or, if you have GPUs, 
# python run.py --multirun dataset=siftsmall,sift1m,deep1m algo=linear,annoy,ivfpq,hnsw,linear_gpu,ivfpq_gpu,ivfpq4bit,scann,pq,pq4bit,hnsw_faiss,nsg

python plot.py

How it works

Config

All config files are on ./conf. You can edit the config files to change parameters.

Download

Download a target dataset by python download.py dataset=DATASET. Several datasets can be downloaded at once by python download.py --multirun dataset=DATASET1,DATASET2,DATASET3. See hydra for more detailed APIs for multirun.
The data will be placed on ./dataset.

Run

Evaluate a target algorithm (ALGO) with a target dataset (DATASET) by python run.py dataset=DATASET algo=ALGO. You can run multiple algorithms on multiple datasets by python run.py --multirun dataset=DATASET1,DATASET2 algo=ALGO1,ALGO2.
Indices (data structures for search) are stored on ./interim. They are reused for each run with different query parameters.
The search result will be on ./output.
By default, we run each algorithm num_trial=10 times and return the average runtime. You can change this by: python run.py num_trial=5

Plot

You can visualize the search result by python plot.py. This script checks ./output and generate figures for each dataset on ./result_img.
In order not to print query parameters, you can set the flag false: python plot.py with_query_param=false.
To change the size of the image: python plot.py width=15 height=10

Log

When running run.py or plot.py, the output files will be on ./log as well. For example with python run.py algo=annoy dataset=siftsmall, the result file will be saved on (1) ./output/siftsmall/annoy/result.yaml and (2) ./log/2020-03-11/22-30-59/0/result.yaml.

Supported datasets

dataset	dim	#base	#query	#train	misc
siftsmall	128	10,000	100	25,000	A toy dataset for hello-world
sift1m	128	1,000,000	10,000	100,000
deep1m	96	1,000,000	10,000	100,000	The first 1M vectors of Deep1B. Hepler scripts

Supported Algorithms

Note that hnsw (hnswlib) is an original implementation by the original authors, and hnsw (faiss) is a re-implemented version by the faiss team.

Add a new algorithm

Write a wrapper class on ./annbench/algo. The class must inherit BaseANN class. See annoy.py for examples.
Update ./annbench/algo/proxy.py
Add the name of the library on requirements.txt.
Add a config file on ./conf/algo.
Make sure the algorithm runs on a single thread

Add a new dataset

Write a wrapper class that inherits BaseDataset on ./annbench/dataset. An simple example is siftsmall.py.
Update ./annbench/dataset/proxy.py.
Add a config file on ./conf/dataset.

Advanced

Evaluation criteria

We followed the standard evaluation criteria in the field of computer vision, Recall@1. See the evaluation function for more details. These functions are from the benchmark scripts of the faiss library.

Index/query parameters

We define a simple guideline to set parameters. An algorithm has to have several index parameters and a single query parameter. For one set of index parameters, one index (data structure) is built. For this index, we run the search by sweeping the query parameter.
For example with ivfpq, let us consider the following index parameters:
```
param_index={"M": 8, "nlist": 100}
```
With these parameters, one index (let us denote ivfpq(M=8, nlist=100)) is created. This index is stored in the disk as M8_nlist100.bin, where the way of naming is defined in the function stringify_index_param. Here, a query parameter is defined as:
```
param_query={"nprobe": [1, 2, 4, 8, 16]}
```
In the search step, the index is read from the disk onto the memory first. Then we run the search five times, with for nprobe in [1, 2, 4, 8, 16]. This creates five results (five pairs of (recall, runtime)). By connecting these results, one polyline is drawn on the final plot.
Note that you must sort the values of the above query parameter. If you forget to sort (e.g., [1, 4, 2, 8, 16]), the final graph would become weird.

Specialization

Index/query parameters for each algorithm is defined in ./conf/algo/. These parameters are used for all datasets by default. If you'd like to specialize parameters for a specific dataset, you can define the specialized version in ./conf/specialized_param/.
For example, the default parameters for ivfpq is defined here, where nlist=100. You can set nlist=1000 for the sift1m dataset by adding a config file here

Dynamic configuration from the command line

In addition to editing the config files, you can override values from the commandline thanks to hydra.

Examples:

Writing index structures on SOMEWHEE/LARGE_HDD/:

python run.py interim=SOMEWHERE/LARGE_HDD/interim

Run the ivfpq algorithm with different query parameters:

python run.py algo=ivfpq dataset=siftsmall param_query.nprobe=[1,5,25]

Reference

ANN Benchmarks

Contributors

Feel free to open a pull request.

@foreverYoungGitHub (hnsw_faiss and nsg)

Author

Yusuke Matsui

matsui528 / annbench