faiss 0.2.0 → 0.2.4

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

data/vendor/faiss/faiss/impl/NNDescent.h

@@ -0,0 +1,154 @@
+/**
+ * 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++ -*-
+
+#pragma once
+
+#include <algorithm>
+#include <mutex>
+#include <queue>
+#include <random>
+#include <unordered_set>
+#include <vector>
+
+#include <omp.h>
+
+#include <faiss/Index.h>
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/platform_macros.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/utils/random.h>
+
+namespace faiss {
+
+/** Implementation of NNDescent which is one of the most popular
+ *  KNN graph building algorithms
+ *
+ * Efficient K-Nearest Neighbor Graph Construction for Generic
+ * Similarity Measures
+ *
+ *  Dong, Wei, Charikar Moses, and Kai Li, WWW 2011
+ *
+ * This implmentation is heavily influenced by the efanna
+ * implementation by Cong Fu and the KGraph library by Wei Dong
+ * (https://github.com/ZJULearning/efanna_graph)
+ * (https://github.com/aaalgo/kgraph)
+ *
+ * The NNDescent object stores only the neighbor link structure,
+ * see IndexNNDescent.h for the full index object.
+ */
+
+struct VisitedTable;
+struct DistanceComputer;
+
+namespace nndescent {
+
+struct Neighbor {
+    int id;
+    float distance;
+    bool flag;
+
+    Neighbor() = default;
+    Neighbor(int id, float distance, bool f)
+            : id(id), distance(distance), flag(f) {}
+
+    inline bool operator<(const Neighbor& other) const {
+        return distance < other.distance;
+    }
+};
+
+struct Nhood {
+    std::mutex lock;
+    std::vector<Neighbor> pool; // candidate pool (a max heap)
+    int M;                      // number of new neighbors to be operated
+
+    std::vector<int> nn_old;  // old neighbors
+    std::vector<int> nn_new;  // new neighbors
+    std::vector<int> rnn_old; // reverse old neighbors
+    std::vector<int> rnn_new; // reverse new neighbors
+
+    Nhood() = default;
+
+    Nhood(int l, int s, std::mt19937& rng, int N);
+
+    Nhood& operator=(const Nhood& other);
+
+    Nhood(const Nhood& other);
+
+    void insert(int id, float dist);
+
+    template <typename C>
+    void join(C callback) const;
+};
+
+} // namespace nndescent
+
+struct NNDescent {
+    using storage_idx_t = int;
+    using idx_t = Index::idx_t;
+
+    using KNNGraph = std::vector<nndescent::Nhood>;
+
+    explicit NNDescent(const int d, const int K);
+
+    ~NNDescent();
+
+    void build(DistanceComputer& qdis, const int n, bool verbose);
+
+    void search(
+            DistanceComputer& qdis,
+            const int topk,
+            idx_t* indices,
+            float* dists,
+            VisitedTable& vt) const;
+
+    void reset();
+
+    /// Initialize the KNN graph randomly
+    void init_graph(DistanceComputer& qdis);
+
+    /// Perform NNDescent algorithm
+    void nndescent(DistanceComputer& qdis, bool verbose);
+
+    /// Perform local join on each node
+    void join(DistanceComputer& qdis);
+
+    /// Sample new neighbors for each node to peform local join later
+    void update();
+
+    /// Sample a small number of points to evaluate the quality of KNNG built
+    void generate_eval_set(
+            DistanceComputer& qdis,
+            std::vector<int>& c,
+            std::vector<std::vector<int>>& v,
+            int N);
+
+    /// Evaluate the quality of KNNG built
+    float eval_recall(
+            std::vector<int>& ctrl_points,
+            std::vector<std::vector<int>>& acc_eval_set);
+
+    bool has_built;
+
+    int K; // K in KNN graph
+    int S; // number of sample neighbors to be updated for each node
+    int R; // size of reverse links, 0 means the reverse links will not be used
+    int L; // size of the candidate pool in building
+    int iter;        // number of iterations to iterate over
+    int search_L;    // size of candidate pool in searching
+    int random_seed; // random seed for generators
+
+    int d; // dimensions
+
+    int ntotal;
+
+    KNNGraph graph;
+    std::vector<int> final_graph;
+};
+
+} // namespace faiss