faiss 0.1.5 → 0.2.2

data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h

@@ -0,0 +1,172 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <stdint.h>
+
+#include <random>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include <faiss/impl/AdditiveQuantizer.h>
+#include <faiss/utils/utils.h>
+
+namespace faiss {
+
+/** Implementation of LSQ/LSQ++ described in the following two papers:
+ *
+ * Revisiting additive quantization
+ * Julieta Martinez, et al. ECCV 2016
+ *
+ * LSQ++: Lower running time and higher recall in multi-codebook quantization
+ * Julieta Martinez, et al. ECCV 2018
+ *
+ * This implementation is mostly translated from the Julia implementations
+ * by Julieta Martinez:
+ * (https://github.com/una-dinosauria/local-search-quantization,
+ *  https://github.com/una-dinosauria/Rayuela.jl)
+ *
+ * The trained codes are stored in `codebooks` which is called
+ * `centroids` in PQ and RQ.
+ */
+
+struct LocalSearchQuantizer : AdditiveQuantizer {
+    size_t K; ///< number of codes per codebook
+
+    size_t train_iters; ///< number of iterations in training
+
+    size_t encode_ils_iters; ///< iterations of local search in encoding
+    size_t train_ils_iters;  ///< iterations of local search in training
+    size_t icm_iters;        ///< number of iterations in icm
+
+    float p;     ///< temperature factor
+    float lambd; ///< regularization factor
+
+    size_t chunk_size; ///< nb of vectors to encode at a time
+
+    int random_seed; ///< seed for random generator
+    size_t nperts;   ///< number of perturbation in each code
+
+    LocalSearchQuantizer(
+            size_t d,      /* dimensionality of the input vectors */
+            size_t M,      /* number of subquantizers */
+            size_t nbits); /* number of bit per subvector index */
+
+    // Train the local search quantizer
+    void train(size_t n, const float* x) override;
+
+    /** Encode a set of vectors
+     *
+     * @param x      vectors to encode, size n * d
+     * @param codes  output codes, size n * code_size
+     */
+    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
+
+    /** Update codebooks given encodings
+     *
+     * @param x      training vectors, size n * d
+     * @param codes  encoded training vectors, size n * M
+     */
+    void update_codebooks(const float* x, const int32_t* codes, size_t n);
+
+    /** Encode vectors given codebooks using iterative conditional mode (icm).
+     *
+     * @param x      vectors to encode, size n * d
+     * @param codes  output codes, size n * M
+     * @param ils_iters number of iterations of iterative local search
+     */
+    void icm_encode(
+            const float* x,
+            int32_t* codes,
+            size_t n,
+            size_t ils_iters,
+            std::mt19937& gen) const;
+
+    void icm_encode_partial(
+            size_t index,
+            const float* x,
+            int32_t* codes,
+            size_t n,
+            const float* binaries,
+            size_t ils_iters,
+            std::mt19937& gen) const;
+
+    void icm_encode_step(
+            const float* unaries,
+            const float* binaries,
+            int32_t* codes,
+            size_t n) const;
+
+    /** Add some perturbation to codebooks
+     *
+     * @param T         temperature of simulated annealing
+     * @param stddev    standard derivations of each dimension in training data
+     */
+    void perturb_codebooks(
+            float T,
+            const std::vector<float>& stddev,
+            std::mt19937& gen);
+
+    /** Add some perturbation to codes
+     *
+     * @param codes codes to be perturbed, size n * M
+     */
+    void perturb_codes(int32_t* codes, size_t n, std::mt19937& gen) const;
+
+    /** Compute binary terms
+     *
+     * @param binaries binary terms, size M * M * K * K
+     */
+    void compute_binary_terms(float* binaries) const;
+
+    /** Compute unary terms
+     *
+     * @param x       vectors to encode, size n * d
+     * @param unaries unary terms, size n * M * K
+     */
+    void compute_unary_terms(const float* x, float* unaries, size_t n) const;
+
+    /** Helper function to compute reconstruction error
+     *
+     * @param x     vectors to encode, size n * d
+     * @param codes encoded codes, size n * M
+     * @param objs  if it is not null, store reconstruction
+                    error of each vector into it, size n
+     */
+    float evaluate(
+            const int32_t* codes,
+            const float* x,
+            size_t n,
+            float* objs = nullptr) const;
+};
+
+/** A helper struct to count consuming time during training.
+ *  It is NOT thread-safe.
+ */
+struct LSQTimer {
+    std::unordered_map<std::string, double> duration;
+    std::unordered_map<std::string, double> t0;
+    std::unordered_map<std::string, bool> started;
+
+    LSQTimer() {
+        reset();
+    }
+
+    double get(const std::string& name);
+
+    void start(const std::string& name);
+
+    void end(const std::string& name);
+
+    void reset();
+};
+
+FAISS_API extern LSQTimer lsq_timer; ///< timer to count consuming time
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/NNDescent.cpp

@@ -0,0 +1,487 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+#include <faiss/impl/NNDescent.h>
+
+#include <mutex>
+#include <string>
+
+#include <faiss/impl/AuxIndexStructures.h>
+
+namespace faiss {
+
+using LockGuard = std::lock_guard<std::mutex>;
+
+namespace nndescent {
+
+void gen_random(std::mt19937& rng, int* addr, const int size, const int N);
+
+Nhood::Nhood(int l, int s, std::mt19937& rng, int N) {
+    M = s;
+    nn_new.resize(s * 2);
+    gen_random(rng, nn_new.data(), (int)nn_new.size(), N);
+}
+
+/// Copy operator
+Nhood& Nhood::operator=(const Nhood& other) {
+    M = other.M;
+    std::copy(
+            other.nn_new.begin(),
+            other.nn_new.end(),
+            std::back_inserter(nn_new));
+    nn_new.reserve(other.nn_new.capacity());
+    pool.reserve(other.pool.capacity());
+    return *this;
+}
+
+/// Copy constructor
+Nhood::Nhood(const Nhood& other) {
+    M = other.M;
+    std::copy(
+            other.nn_new.begin(),
+            other.nn_new.end(),
+            std::back_inserter(nn_new));
+    nn_new.reserve(other.nn_new.capacity());
+    pool.reserve(other.pool.capacity());
+}
+
+/// Insert a point into the candidate pool
+void Nhood::insert(int id, float dist) {
+    LockGuard guard(lock);
+    if (dist > pool.front().distance)
+        return;
+    for (int i = 0; i < pool.size(); i++) {
+        if (id == pool[i].id)
+            return;
+    }
+    if (pool.size() < pool.capacity()) {
+        pool.push_back(Neighbor(id, dist, true));
+        std::push_heap(pool.begin(), pool.end());
+    } else {
+        std::pop_heap(pool.begin(), pool.end());
+        pool[pool.size() - 1] = Neighbor(id, dist, true);
+        std::push_heap(pool.begin(), pool.end());
+    }
+}
+
+/// In local join, two objects are compared only if at least
+/// one of them is new.
+template <typename C>
+void Nhood::join(C callback) const {
+    for (int const i : nn_new) {
+        for (int const j : nn_new) {
+            if (i < j) {
+                callback(i, j);
+            }
+        }
+        for (int j : nn_old) {
+            callback(i, j);
+        }
+    }
+}
+
+void gen_random(std::mt19937& rng, int* addr, const int size, const int N) {
+    for (int i = 0; i < size; ++i) {
+        addr[i] = rng() % (N - size);
+    }
+    std::sort(addr, addr + size);
+    for (int i = 1; i < size; ++i) {
+        if (addr[i] <= addr[i - 1]) {
+            addr[i] = addr[i - 1] + 1;
+        }
+    }
+    int off = rng() % N;
+    for (int i = 0; i < size; ++i) {
+        addr[i] = (addr[i] + off) % N;
+    }
+}
+
+// Insert a new point into the candidate pool in ascending order
+int insert_into_pool(Neighbor* addr, int size, Neighbor nn) {
+    // find the location to insert
+    int left = 0, right = size - 1;
+    if (addr[left].distance > nn.distance) {
+        memmove((char*)&addr[left + 1], &addr[left], size * sizeof(Neighbor));
+        addr[left] = nn;
+        return left;
+    }
+    if (addr[right].distance < nn.distance) {
+        addr[size] = nn;
+        return size;
+    }
+    while (left < right - 1) {
+        int mid = (left + right) / 2;
+        if (addr[mid].distance > nn.distance)
+            right = mid;
+        else
+            left = mid;
+    }
+    // check equal ID
+
+    while (left > 0) {
+        if (addr[left].distance < nn.distance)
+            break;
+        if (addr[left].id == nn.id)
+            return size + 1;
+        left--;
+    }
+    if (addr[left].id == nn.id || addr[right].id == nn.id)
+        return size + 1;
+    memmove((char*)&addr[right + 1],
+            &addr[right],
+            (size - right) * sizeof(Neighbor));
+    addr[right] = nn;
+    return right;
+}
+
+} // namespace nndescent
+
+using namespace nndescent;
+
+constexpr int NUM_EVAL_POINTS = 100;
+
+NNDescent::NNDescent(const int d, const int K) : K(K), random_seed(2021), d(d) {
+    ntotal = 0;
+    has_built = false;
+    S = 10;
+    R = 100;
+    L = K + 50;
+    iter = 10;
+    search_L = 0;
+}
+
+NNDescent::~NNDescent() {}
+
+void NNDescent::join(DistanceComputer& qdis) {
+#pragma omp parallel for default(shared) schedule(dynamic, 100)
+    for (int n = 0; n < ntotal; n++) {
+        graph[n].join([&](int i, int j) {
+            if (i != j) {
+                float dist = qdis.symmetric_dis(i, j);
+                graph[i].insert(j, dist);
+                graph[j].insert(i, dist);
+            }
+        });
+    }
+}
+
+/// Sample neighbors for each node to peform local join later
+/// Store them in nn_new and nn_old
+void NNDescent::update() {
+    // Step 1.
+    // Clear all nn_new and nn_old
+#pragma omp parallel for
+    for (int i = 0; i < ntotal; i++) {
+        std::vector<int>().swap(graph[i].nn_new);
+        std::vector<int>().swap(graph[i].nn_old);
+    }
+
+    // Step 2.
+    // Compute the number of neighbors which is new i.e. flag is true
+    // in the candidate pool. This must not exceed the sample number S.
+    // That means We only select S new neighbors.
+#pragma omp parallel for
+    for (int n = 0; n < ntotal; ++n) {
+        auto& nn = graph[n];
+        std::sort(nn.pool.begin(), nn.pool.end());
+
+        if (nn.pool.size() > L)
+            nn.pool.resize(L);
+        nn.pool.reserve(L); // keep the pool size be L
+
+        int maxl = std::min(nn.M + S, (int)nn.pool.size());
+        int c = 0;
+        int l = 0;
+
+        while ((l < maxl) && (c < S)) {
+            if (nn.pool[l].flag)
+                ++c;
+            ++l;
+        }
+        nn.M = l;
+    }
+
+    // Step 3.
+    // Find reverse links for each node
+    // Randomly choose R reverse links.
+#pragma omp parallel
+    {
+        std::mt19937 rng(random_seed * 5081 + omp_get_thread_num());
+#pragma omp for
+        for (int n = 0; n < ntotal; ++n) {
+            auto& node = graph[n];
+            auto& nn_new = node.nn_new;
+            auto& nn_old = node.nn_old;
+
+            for (int l = 0; l < node.M; ++l) {
+                auto& nn = node.pool[l];
+                auto& other = graph[nn.id]; // the other side of the edge
+
+                if (nn.flag) { // the node is inserted newly
+                    // push the neighbor into nn_new
+                    nn_new.push_back(nn.id);
+                    // push itself into other.rnn_new if it is not in
+                    // the candidate pool of the other side
+                    if (nn.distance > other.pool.back().distance) {
+                        LockGuard guard(other.lock);
+                        if (other.rnn_new.size() < R) {
+                            other.rnn_new.push_back(n);
+                        } else {
+                            int pos = rng() % R;
+                            other.rnn_new[pos] = n;
+                        }
+                    }
+                    nn.flag = false;
+
+                } else { // the node is old
+                    // push the neighbor into nn_old
+                    nn_old.push_back(nn.id);
+                    // push itself into other.rnn_old if it is not in
+                    // the candidate pool of the other side
+                    if (nn.distance > other.pool.back().distance) {
+                        LockGuard guard(other.lock);
+                        if (other.rnn_old.size() < R) {
+                            other.rnn_old.push_back(n);
+                        } else {
+                            int pos = rng() % R;
+                            other.rnn_old[pos] = n;
+                        }
+                    }
+                }
+            }
+            // make heap to join later (in join() function)
+            std::make_heap(node.pool.begin(), node.pool.end());
+        }
+    }
+
+    // Step 4.
+    // Combine the forward and the reverse links
+    // R = 0 means no reverse links are used.
+#pragma omp parallel for
+    for (int i = 0; i < ntotal; ++i) {
+        auto& nn_new = graph[i].nn_new;
+        auto& nn_old = graph[i].nn_old;
+        auto& rnn_new = graph[i].rnn_new;
+        auto& rnn_old = graph[i].rnn_old;
+
+        nn_new.insert(nn_new.end(), rnn_new.begin(), rnn_new.end());
+        nn_old.insert(nn_old.end(), rnn_old.begin(), rnn_old.end());
+        if (nn_old.size() > R * 2) {
+            nn_old.resize(R * 2);
+            nn_old.reserve(R * 2);
+        }
+
+        std::vector<int>().swap(graph[i].rnn_new);
+        std::vector<int>().swap(graph[i].rnn_old);
+    }
+}
+
+void NNDescent::nndescent(DistanceComputer& qdis, bool verbose) {
+    int num_eval_points = std::min(NUM_EVAL_POINTS, ntotal);
+    std::vector<int> eval_points(num_eval_points);
+    std::vector<std::vector<int>> acc_eval_set(num_eval_points);
+    std::mt19937 rng(random_seed * 6577 + omp_get_thread_num());
+    gen_random(rng, eval_points.data(), eval_points.size(), ntotal);
+    generate_eval_set(qdis, eval_points, acc_eval_set, ntotal);
+    for (int it = 0; it < iter; it++) {
+        join(qdis);
+        update();
+
+        if (verbose) {
+            float recall = eval_recall(eval_points, acc_eval_set);
+            printf("Iter: %d, recall@%d: %lf\n", it, K, recall);
+        }
+    }
+}
+
+/// Sample a small number of points to evaluate the quality of KNNG built
+void NNDescent::generate_eval_set(
+        DistanceComputer& qdis,
+        std::vector<int>& c,
+        std::vector<std::vector<int>>& v,
+        int N) {
+#pragma omp parallel for
+    for (int i = 0; i < c.size(); i++) {
+        std::vector<Neighbor> tmp;
+        for (int j = 0; j < N; j++) {
+            if (i == j)
+                continue; // skip itself
+            float dist = qdis.symmetric_dis(c[i], j);
+            tmp.push_back(Neighbor(j, dist, true));
+        }
+
+        std::partial_sort(tmp.begin(), tmp.begin() + K, tmp.end());
+        for (int j = 0; j < K; j++) {
+            v[i].push_back(tmp[j].id);
+        }
+    }
+}
+
+/// Evaluate the quality of KNNG built
+float NNDescent::eval_recall(
+        std::vector<int>& eval_points,
+        std::vector<std::vector<int>>& acc_eval_set) {
+    float mean_acc = 0.0f;
+    for (size_t i = 0; i < eval_points.size(); i++) {
+        float acc = 0;
+        std::vector<Neighbor>& g = graph[eval_points[i]].pool;
+        std::vector<int>& v = acc_eval_set[i];
+        for (size_t j = 0; j < g.size(); j++) {
+            for (size_t k = 0; k < v.size(); k++) {
+                if (g[j].id == v[k]) {
+                    acc++;
+                    break;
+                }
+            }
+        }
+        mean_acc += acc / v.size();
+    }
+    return mean_acc / eval_points.size();
+}
+
+/// Initialize the KNN graph randomly
+void NNDescent::init_graph(DistanceComputer& qdis) {
+    graph.reserve(ntotal);
+    {
+        std::mt19937 rng(random_seed * 6007);
+        for (int i = 0; i < ntotal; i++) {
+            graph.push_back(Nhood(L, S, rng, (int)ntotal));
+        }
+    }
+#pragma omp parallel
+    {
+        std::mt19937 rng(random_seed * 7741 + omp_get_thread_num());
+#pragma omp for
+        for (int i = 0; i < ntotal; i++) {
+            std::vector<int> tmp(S);
+
+            gen_random(rng, tmp.data(), S, ntotal);
+
+            for (int j = 0; j < S; j++) {
+                int id = tmp[j];
+                if (id == i)
+                    continue;
+                float dist = qdis.symmetric_dis(i, id);
+
+                graph[i].pool.push_back(Neighbor(id, dist, true));
+            }
+            std::make_heap(graph[i].pool.begin(), graph[i].pool.end());
+            graph[i].pool.reserve(L);
+        }
+    }
+}
+
+void NNDescent::build(DistanceComputer& qdis, const int n, bool verbose) {
+    FAISS_THROW_IF_NOT_MSG(L >= K, "L should be >= K in NNDescent.build");
+
+    if (verbose) {
+        printf("Parameters: K=%d, S=%d, R=%d, L=%d, iter=%d\n",
+               K,
+               S,
+               R,
+               L,
+               iter);
+    }
+
+    ntotal = n;
+    init_graph(qdis);
+    nndescent(qdis, verbose);
+
+    final_graph.resize(ntotal * K);
+
+    // Store the neighbor link structure into final_graph
+    // Clear the old graph
+    for (int i = 0; i < ntotal; i++) {
+        std::sort(graph[i].pool.begin(), graph[i].pool.end());
+        for (int j = 0; j < K; j++) {
+            FAISS_ASSERT(graph[i].pool[j].id < ntotal);
+            final_graph[i * K + j] = graph[i].pool[j].id;
+        }
+    }
+    std::vector<Nhood>().swap(graph);
+    has_built = true;
+
+    if (verbose) {
+        printf("Addes %d points into the index\n", ntotal);
+    }
+}
+
+void NNDescent::search(
+        DistanceComputer& qdis,
+        const int topk,
+        idx_t* indices,
+        float* dists,
+        VisitedTable& vt) const {
+    FAISS_THROW_IF_NOT_MSG(has_built, "The index is not build yet.");
+    int L = std::max(search_L, topk);
+
+    // candidate pool, the K best items is the result.
+    std::vector<Neighbor> retset(L + 1);
+
+    // Randomly choose L points to intialize the candidate pool
+    std::vector<int> init_ids(L);
+    std::mt19937 rng(random_seed);
+
+    gen_random(rng, init_ids.data(), L, ntotal);
+    for (int i = 0; i < L; i++) {
+        int id = init_ids[i];
+        float dist = qdis(id);
+        retset[i] = Neighbor(id, dist, true);
+    }
+
+    // Maintain the candidate pool in ascending order
+    std::sort(retset.begin(), retset.begin() + L);
+
+    int k = 0;
+
+    // Stop until the smallest position updated is >= L
+    while (k < L) {
+        int nk = L;
+
+        if (retset[k].flag) {
+            retset[k].flag = false;
+            int n = retset[k].id;
+
+            for (int m = 0; m < K; ++m) {
+                int id = final_graph[n * K + m];
+                if (vt.get(id))
+                    continue;
+
+                vt.set(id);
+                float dist = qdis(id);
+                if (dist >= retset[L - 1].distance)
+                    continue;
+
+                Neighbor nn(id, dist, true);
+                int r = insert_into_pool(retset.data(), L, nn);
+
+                if (r < nk)
+                    nk = r;
+            }
+        }
+        if (nk <= k)
+            k = nk;
+        else
+            ++k;
+    }
+    for (size_t i = 0; i < topk; i++) {
+        indices[i] = retset[i].id;
+        dists[i] = retset[i].distance;
+    }
+
+    vt.advance();
+}
+
+void NNDescent::reset() {
+    has_built = false;
+    ntotal = 0;
+    final_graph.resize(0);
+}
+
+} // namespace faiss