Initial commit

tests/demo_imi_flat.cpp

@@ -0,0 +1,154 @@
+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+
+
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+
+#include <sys/time.h>
+
+
+#include <faiss/IndexPQ.h>
+#include <faiss/IndexIVF.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/index_io.h>
+
+double elapsed ()
+{
+    struct timeval tv;
+    gettimeofday (&tv, nullptr);
+    return  tv.tv_sec + tv.tv_usec * 1e-6;
+}
+
+
+int main ()
+{
+    double t0 = elapsed();
+
+    // dimension of the vectors to index
+    int d = 128;
+
+    // size of the database we plan to index
+    size_t nb = 1000 * 1000;
+
+    // make a set of nt training vectors in the unit cube
+    // (could be the database)
+    size_t nt = 100 * 1000;
+
+    //---------------------------------------------------------------
+    // Define the core quantizer
+    // We choose a multiple inverted index for faster training with less data
+    // and because it usually offers best accuracy/speed trade-offs
+    //
+    // We here assume that its lifespan of this coarse quantizer will cover the
+    // lifespan of the inverted-file quantizer IndexIVFFlat below
+    // With dynamic allocation, one may give the responsability to free the
+    // quantizer to the inverted-file index (with attribute do_delete_quantizer)
+    //
+    // Note: a regular clustering algorithm would be defined as:
+    //       faiss::IndexFlatL2 coarse_quantizer (d);
+    //
+    // Use nhash=2 subquantizers used to define the product coarse quantizer
+    // Number of bits: we will have 2^nbits_coarse centroids per subquantizer
+    //                 meaning (2^12)^nhash distinct inverted lists
+    size_t nhash = 2;
+    size_t nbits_subq = int (log2 (nb+1) / 2);        // good choice in general
+    size_t ncentroids = 1 << (nhash * nbits_subq);  // total # of centroids
+
+    faiss::MultiIndexQuantizer coarse_quantizer (d, nhash, nbits_subq);
+
+    printf ("IMI (%ld,%ld): %ld virtual centroids (target: %ld base vectors)",
+            nhash, nbits_subq, ncentroids, nb);
+
+    // the coarse quantizer should not be dealloced before the index
+    // 4 = nb of bytes per code (d must be a multiple of this)
+    // 8 = nb of bits per sub-code (almost always 8)
+    faiss::MetricType metric = faiss::METRIC_L2; // can be METRIC_INNER_PRODUCT
+    faiss::IndexIVFFlat index (&coarse_quantizer, d, ncentroids, metric);
+    index.quantizer_trains_alone = true;
+
+    // define the number of probes. 2048 is for high-dim, overkilled in practice
+    // Use 4-1024 depending on the trade-off speed accuracy that you want
+    index.nprobe = 2048;
+
+
+    { // training
+        printf ("[%.3f s] Generating %ld vectors in %dD for training\n",
+                elapsed() - t0, nt, d);
+
+        std::vector <float> trainvecs (nt * d);
+        for (size_t i = 0; i < nt * d; i++) {
+            trainvecs[i] = drand48();
+        }
+
+        printf ("[%.3f s] Training the index\n", elapsed() - t0);
+        index.verbose = true;
+        index.train (nt, trainvecs.data());
+    }
+
+    size_t nq;
+    std::vector<float> queries;
+
+    { // populating the database
+        printf ("[%.3f s] Building a dataset of %ld vectors to index\n",
+                elapsed() - t0, nb);
+
+        std::vector <float> database (nb * d);
+        for (size_t i = 0; i < nb * d; i++) {
+            database[i] = drand48();
+        }
+
+        printf ("[%.3f s] Adding the vectors to the index\n", elapsed() - t0);
+
+        index.add (nb, database.data());
+
+        // remember a few elements from the database as queries
+        int i0 = 1234;
+        int i1 = 1244;
+
+        nq = i1 - i0;
+        queries.resize (nq * d);
+        for (int i = i0; i < i1; i++) {
+            for (int j = 0; j < d; j++) {
+                queries [(i - i0) * d  + j]  = database [i * d + j];
+            }
+        }
+    }
+
+    { // searching the database
+        int k = 5;
+        printf ("[%.3f s] Searching the %d nearest neighbors "
+                "of %ld vectors in the index\n",
+                elapsed() - t0, k, nq);
+
+        std::vector<faiss::Index::idx_t> nns (k * nq);
+        std::vector<float>               dis (k * nq);
+
+        index.search (nq, queries.data(), k, dis.data(), nns.data());
+
+        printf ("[%.3f s] Query results (vector ids, then distances):\n",
+                elapsed() - t0);
+
+        for (int i = 0; i < nq; i++) {
+            printf ("query %2d: ", i);
+            for (int j = 0; j < k; j++) {
+                printf ("%7ld ", nns[j + i * k]);
+            }
+            printf ("\n     dis: ");
+            for (int j = 0; j < k; j++) {
+                printf ("%7g ", dis[j + i * k]);
+            }
+            printf ("\n");
+        }
+    }
+    return 0;
+}

tests/demo_ivfpq_indexing.cpp

@@ -0,0 +1,149 @@
+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+
+
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+
+#include <sys/time.h>
+
+
+#include "../IndexIVFPQ.h"
+#include "../IndexFlat.h"
+#include "../index_io.h"
+
+double elapsed ()
+{
+    struct timeval tv;
+    gettimeofday (&tv, NULL);
+    return  tv.tv_sec + tv.tv_usec * 1e-6;
+}
+
+
+int main ()
+{
+
+    double t0 = elapsed();
+
+    // dimension of the vectors to index
+    int d = 128;
+
+    // size of the database we plan to index
+    size_t nb = 200 * 1000;
+
+    // make a set of nt training vectors in the unit cube
+    // (could be the database)
+    size_t nt = 100 * 1000;
+
+    // make the index object and train it
+    faiss::IndexFlatL2 coarse_quantizer (d);
+
+    // a reasonable number of centroids to index nb vectors
+    int ncentroids = int (4 * sqrt (nb));
+
+    // the coarse quantizer should not be dealloced before the index
+    // 4 = nb of bytes per code (d must be a multiple of this)
+    // 8 = nb of bits per sub-code (almost always 8)
+    faiss::IndexIVFPQ index (&coarse_quantizer, d,
+                             ncentroids, 4, 8);
+
+
+    { // training
+        printf ("[%.3f s] Generating %ld vectors in %dD for training\n",
+                elapsed() - t0, nt, d);
+
+        std::vector <float> trainvecs (nt * d);
+        for (size_t i = 0; i < nt * d; i++) {
+            trainvecs[i] = drand48();
+        }
+
+        printf ("[%.3f s] Training the index\n",
+                elapsed() - t0);
+        index.verbose = true;
+
+        index.train (nt, trainvecs.data());
+    }
+
+    { // I/O demo
+        const char *outfilename = "/tmp/index_trained.faissindex";
+        printf ("[%.3f s] storing the pre-trained index to %s\n",
+                elapsed() - t0, outfilename);
+
+        write_index (&index, outfilename);
+    }
+
+    size_t nq;
+    std::vector<float> queries;
+
+    { // populating the database
+        printf ("[%.3f s] Building a dataset of %ld vectors to index\n",
+                elapsed() - t0, nb);
+
+        std::vector <float> database (nb * d);
+        for (size_t i = 0; i < nb * d; i++) {
+            database[i] = drand48();
+        }
+
+        printf ("[%.3f s] Adding the vectors to the index\n",
+                elapsed() - t0);
+
+        index.add (nb, database.data());
+
+        printf ("[%.3f s] imbalance factor: %g\n",
+                elapsed() - t0, index.imbalance_factor ());
+
+        // remember a few elements from the database as queries
+        int i0 = 1234;
+        int i1 = 1243;
+
+        nq = i1 - i0;
+        queries.resize (nq * d);
+        for (int i = i0; i < i1; i++) {
+            for (int j = 0; j < d; j++) {
+                queries [(i - i0) * d  + j]  = database [i * d + j];
+            }
+        }
+
+    }
+
+    { // searching the database
+        int k = 5;
+        printf ("[%.3f s] Searching the %d nearest neighbors "
+                "of %ld vectors in the index\n",
+                elapsed() - t0, k, nq);
+
+        std::vector<faiss::Index::idx_t> nns (k * nq);
+        std::vector<float>               dis (k * nq);
+
+        index.search (nq, queries.data(), k, dis.data(), nns.data());
+
+        printf ("[%.3f s] Query results (vector ids, then distances):\n",
+                elapsed() - t0);
+
+        for (int i = 0; i < nq; i++) {
+            printf ("query %2d: ", i);
+            for (int j = 0; j < k; j++) {
+                printf ("%7ld ", nns[j + i * k]);
+            }
+            printf ("\n     dis: ");
+            for (int j = 0; j < k; j++) {
+                printf ("%7g ", dis[j + i * k]);
+            }
+            printf ("\n");
+        }
+
+        printf ("note that the nearest neighbor is not at "
+                "distance 0 due to quantization errors\n");
+    }
+
+    return 0;
+}

tests/demo_sift1M.cpp

@@ -0,0 +1,255 @@
+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+ // Copyright 2004-present Facebook. All Rights Reserved
+
+
+#include <cmath>
+#include <cstdio>
+#include <cstdlib>
+#include <cassert>
+#include <cstring>
+
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <sys/time.h>
+
+#include "../AutoTune.h"
+
+
+/**
+ * To run this demo, please download the ANN_SIFT1M dataset from
+ *
+ *   http://corpus-texmex.irisa.fr/
+ *
+ * and unzip it to the sudirectory sift1M.
+ **/
+
+/*****************************************************
+ * I/O functions for fvecs and ivecs
+ *****************************************************/
+
+
+float * fvecs_read (const char *fname,
+                    size_t *d_out, size_t *n_out)
+{
+    FILE *f = fopen(fname, "r");
+    if(!f) {
+        fprintf(stderr, "could not open %s\n", fname);
+        perror("");
+        abort();
+    }
+    int d;
+    fread(&d, 1, sizeof(int), f);
+    assert((d > 0 && d < 1000000) || !"unreasonable dimension");
+    fseek(f, 0, SEEK_SET);
+    struct stat st;
+    fstat(fileno(f), &st);
+    size_t sz = st.st_size;
+    assert(sz % ((d + 1) * 4) == 0 || !"weird file size");
+    size_t n = sz / ((d + 1) * 4);
+
+    *d_out = d; *n_out = n;
+    float *x = new float[n * (d + 1)];
+    size_t nr = fread(x, sizeof(float), n * (d + 1), f);
+    assert(nr == n * (d + 1) || !"could not read whole file");
+
+    // shift array to remove row headers
+    for(size_t i = 0; i < n; i++)
+        memmove(x + i * d, x + 1 + i * (d + 1), d * sizeof(*x));
+
+    fclose(f);
+    return x;
+}
+
+// not very clean, but works as long as sizeof(int) == sizeof(float)
+int *ivecs_read(const char *fname, size_t *d_out, size_t *n_out)
+{
+    return (int*)fvecs_read(fname, d_out, n_out);
+}
+
+double elapsed ()
+{
+    struct timeval tv;
+    gettimeofday (&tv, nullptr);
+    return  tv.tv_sec + tv.tv_usec * 1e-6;
+}
+
+
+
+int main()
+{
+    double t0 = elapsed();
+
+    // this is typically the fastest one.
+    const char *index_key = "IVF4096,Flat";
+
+    // these ones have better memory usage
+    // const char *index_key = "Flat";
+    // const char *index_key = "PQ32";
+    // const char *index_key = "PCA80,Flat";
+    // const char *index_key = "IVF4096,PQ8+16";
+    // const char *index_key = "IVF4096,PQ32";
+    // const char *index_key = "IMI2x8,PQ32";
+    // const char *index_key = "IMI2x8,PQ8+16";
+    // const char *index_key = "OPQ16_64,IMI2x8,PQ8+16";
+
+    faiss::Index * index;
+
+    size_t d;
+
+    {
+        printf ("[%.3f s] Loading train set\n", elapsed() - t0);
+
+        size_t nt;
+        float *xt = fvecs_read("sift1M/sift_learn.fvecs", &d, &nt);
+
+        printf ("[%.3f s] Preparing index \"%s\" d=%ld\n",
+                elapsed() - t0, index_key, d);
+        index = faiss::index_factory(d, index_key);
+
+        printf ("[%.3f s] Training on %ld vectors\n", elapsed() - t0, nt);
+
+        index->train(nt, xt);
+        delete [] xt;
+    }
+
+
+    {
+        printf ("[%.3f s] Loading database\n", elapsed() - t0);
+
+        size_t nb, d2;
+        float *xb = fvecs_read("sift1M/sift_base.fvecs", &d2, &nb);
+        assert(d == d2 || !"dataset does not have same dimension as train set");
+
+        printf ("[%.3f s] Indexing database, size %ld*%ld\n",
+                elapsed() - t0, nb, d);
+
+        index->add(nb, xb);
+
+        delete [] xb;
+    }
+
+    size_t nq;
+    float *xq;
+
+    {
+        printf ("[%.3f s] Loading queries\n", elapsed() - t0);
+
+        size_t d2;
+        xq = fvecs_read("sift1M/sift_query.fvecs", &d2, &nq);
+        assert(d == d2 || !"query does not have same dimension as train set");
+
+    }
+
+    size_t k; // nb of results per query in the GT
+    faiss::Index::idx_t *gt;  // nq * k matrix of ground-truth nearest-neighbors
+
+    {
+        printf ("[%.3f s] Loading ground truth for %ld queries\n",
+                elapsed() - t0, nq);
+
+        // load ground-truth and convert int to long
+        size_t nq2;
+        int *gt_int = ivecs_read("sift1M/sift_groundtruth.ivecs", &k, &nq2);
+        assert(nq2 == nq || !"incorrect nb of ground truth entries");
+
+        gt = new faiss::Index::idx_t[k * nq];
+        for(int i = 0; i < k * nq; i++) {
+            gt[i] = gt_int[i];
+        }
+        delete [] gt_int;
+    }
+
+    // Result of the auto-tuning
+    std::string selected_params;
+
+    { // run auto-tuning
+
+        printf ("[%.3f s] Preparing auto-tune criterion 1-recall at 1 "
+                "criterion, with k=%ld nq=%ld\n", elapsed() - t0, k, nq);
+
+        faiss::OneRecallAtRCriterion crit(nq, 1);
+        crit.set_groundtruth (k, nullptr, gt);
+        crit.nnn = k; // by default, the criterion will request only 1 NN
+
+        printf ("[%.3f s] Preparing auto-tune parameters\n", elapsed() - t0);
+
+        faiss::ParameterSpace params;
+        params.initialize(index);
+
+        printf ("[%.3f s] Auto-tuning over %ld parameters (%ld combinations)\n",
+                elapsed() - t0, params.parameter_ranges.size(),
+                params.n_combinations());
+
+        faiss::OperatingPoints ops;
+        params.explore (index, nq, xq, crit, &ops);
+
+        printf ("[%.3f s] Found the following operating points: \n",
+                elapsed() - t0);
+
+        ops.display ();
+
+        // keep the first parameter that obtains > 0.5 1-recall@1
+        for (int i = 0; i < ops.optimal_pts.size(); i++) {
+            if (ops.optimal_pts[i].perf > 0.5) {
+                selected_params = ops.optimal_pts[i].key;
+                break;
+            }
+        }
+        assert (selected_params.size() >= 0 ||
+                !"could not find good enough op point");
+    }
+
+
+    { // Use the found configuration to perform a search
+
+        faiss::ParameterSpace params;
+
+        printf ("[%.3f s] Setting parameter configuration \"%s\" on index\n",
+                elapsed() - t0, selected_params.c_str());
+
+        params.set_index_parameters (index, selected_params.c_str());
+
+        printf ("[%.3f s] Perform a search on %ld queries\n",
+                elapsed() - t0, nq);
+
+        // output buffers
+        faiss::Index::idx_t *I = new  faiss::Index::idx_t[nq * k];
+        float *D = new float[nq * k];
+
+        index->search(nq, xq, k, D, I);
+
+        printf ("[%.3f s] Compute recalls\n", elapsed() - t0);
+
+        // evaluate result by hand.
+        int n_1 = 0, n_10 = 0, n_100 = 0;
+        for(int i = 0; i < nq; i++) {
+            int gt_nn = gt[i * k];
+            for(int j = 0; j < k; j++) {
+                if (I[i * k + j] == gt_nn) {
+                    if(j < 1) n_1++;
+                    if(j < 10) n_10++;
+                    if(j < 100) n_100++;
+                }
+            }
+        }
+        printf("R@1 = %.4f\n", n_1 / float(nq));
+        printf("R@10 = %.4f\n", n_10 / float(nq));
+        printf("R@100 = %.4f\n", n_100 / float(nq));
+
+    }
+
+    delete [] xq;
+    delete [] gt;
+    delete index;
+    return 0;
+}

tests/test_blas.cpp

@@ -0,0 +1,86 @@
+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <cstdio>
+#include <cstdlib>
+
+#undef FINTEGER
+#define FINTEGER long
+
+
+extern "C" {
+
+/* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
+
+int sgemm_ (const char *transa, const char *transb, FINTEGER *m, FINTEGER *
+            n, FINTEGER *k, const float *alpha, const float *a,
+            FINTEGER *lda, const float *b, FINTEGER *
+            ldb, float *beta, float *c, FINTEGER *ldc);
+
+/* Lapack functions, see http://www.netlib.org/clapack/old/single/sgeqrf.c */
+
+int sgeqrf_ (FINTEGER *m, FINTEGER *n, float *a, FINTEGER *lda,
+                 float *tau, float *work, FINTEGER *lwork, FINTEGER *info);
+
+}
+
+float *new_random_vec(int size)
+{
+    float *x = new float[size];
+    for (int i = 0; i < size; i++)
+        x[i] = drand48();
+    return x;
+}
+
+
+int main() {
+
+    FINTEGER m = 10, n = 20, k = 30;
+    float *a = new_random_vec(m * k), *b = new_random_vec(n * k), *c = new float[n * m];
+    float one = 1.0, zero = 0.0;
+
+    printf("BLAS test\n");
+
+    sgemm_("Not transposed", "Not transposed",
+           &m, &n, &k, &one, a, &m, b, &k, &zero, c, &m);
+
+    printf("errors=\n");
+
+    for (int i = 0; i < m; i++) {
+        for (int j = 0; j < n; j++) {
+            float accu = 0;
+            for (int l = 0; l < k; l++)
+                accu += a[i + l * m] * b[l + j * k];
+            printf ("%6.3f ", accu - c[i + j * m]);
+        }
+        printf("\n");
+    }
+
+    long info = 0x64bL << 32;
+    long mi = 0x64bL << 32 | m;
+    float *tau = new float[m];
+    FINTEGER lwork = -1;
+
+    float work1;
+
+    printf("Intentional Lapack error (appears only for 64-bit INTEGER):\n");
+    sgeqrf_ (&mi, &n, c, &m, tau, &work1, &lwork, (FINTEGER*)&info);
+
+    // sgeqrf_ (&m, &n, c, &zeroi, tau, &work1, &lwork, (FINTEGER*)&info);
+    printf("info=%016lx\n", info);
+
+    if(info >> 32 == 0x64b) {
+        printf("Lapack uses 32-bit integers\n");
+    } else {
+        printf("Lapack uses 64-bit integers\n");
+    }
+
+
+    return 0;
+}

tests/test_ivfpq_indexing.cpp

@@ -0,0 +1,101 @@
+
+/**
+ * Copyright (c) 2015-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the CC-by-NC license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// Copyright 2004-present Facebook. All Rights Reserved
+
+#include <cstdio>
+#include <cstdlib>
+
+#include <gtest/gtest.h>
+
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/index_io.h>
+
+TEST(IVFPQ, accuracy) {
+
+    // dimension of the vectors to index
+    int d = 64;
+
+    // size of the database we plan to index
+    size_t nb = 1000;
+
+    // make a set of nt training vectors in the unit cube
+    // (could be the database)
+    size_t nt = 1500;
+
+    // make the index object and train it
+    faiss::IndexFlatL2 coarse_quantizer (d);
+
+    // a reasonable number of cetroids to index nb vectors
+    int ncentroids = 25;
+
+    faiss::IndexIVFPQ index (&coarse_quantizer, d,
+                             ncentroids, 16, 8);
+
+    // index that gives the ground-truth
+    faiss::IndexFlatL2 index_gt (d);
+
+    srand48 (35);
+
+    { // training
+
+        std::vector <float> trainvecs (nt * d);
+        for (size_t i = 0; i < nt * d; i++) {
+            trainvecs[i] = drand48();
+        }
+        index.verbose = true;
+        index.train (nt, trainvecs.data());
+    }
+
+    { // populating the database
+
+        std::vector <float> database (nb * d);
+        for (size_t i = 0; i < nb * d; i++) {
+            database[i] = drand48();
+        }
+
+        index.add (nb, database.data());
+        index_gt.add (nb, database.data());
+    }
+
+    int nq = 200;
+    int n_ok;
+
+    { // searching the database
+
+        std::vector <float> queries (nq * d);
+        for (size_t i = 0; i < nq * d; i++) {
+            queries[i] = drand48();
+        }
+
+        std::vector<faiss::Index::idx_t> gt_nns (nq);
+        std::vector<float>               gt_dis (nq);
+
+        index_gt.search (nq, queries.data(), 1,
+                         gt_dis.data(), gt_nns.data());
+
+        index.nprobe = 5;
+        int k = 5;
+        std::vector<faiss::Index::idx_t> nns (k * nq);
+        std::vector<float>               dis (k * nq);
+
+        index.search (nq, queries.data(), k, dis.data(), nns.data());
+
+        n_ok = 0;
+        for (int q = 0; q < nq; q++) {
+
+            for (int i = 0; i < k; i++)
+                if (nns[q * k + i] == gt_nns[q])
+                    n_ok++;
+        }
+        EXPECT_GT(n_ok, nq * 0.4);
+    }
+
+}