faiss 0.1.0 → 0.1.1

data/vendor/faiss/IndexPQ.h

@@ -0,0 +1,199 @@
+/**
+ * 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++ -*-
+
+#ifndef FAISS_INDEX_PQ_H
+#define FAISS_INDEX_PQ_H
+
+#include <stdint.h>
+
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/impl/ProductQuantizer.h>
+#include <faiss/impl/PolysemousTraining.h>
+
+namespace faiss {
+
+
+/** Index based on a product quantizer. Stored vectors are
+ * approximated by PQ codes. */
+struct IndexPQ: Index {
+
+    /// The product quantizer used to encode the vectors
+    ProductQuantizer pq;
+
+    /// Codes. Size ntotal * pq.code_size
+    std::vector<uint8_t> codes;
+
+    /** Constructor.
+     *
+     * @param d      dimensionality of the input vectors
+     * @param M      number of subquantizers
+     * @param nbits  number of bit per subvector index
+     */
+    IndexPQ (int d,                    ///< dimensionality of the input vectors
+             size_t M,                 ///< number of subquantizers
+             size_t nbits,             ///< number of bit per subvector index
+             MetricType metric = METRIC_L2);
+
+    IndexPQ ();
+
+    void train(idx_t n, const float* x) override;
+
+    void add(idx_t n, const float* x) override;
+
+    void search(
+        idx_t n,
+        const float* x,
+        idx_t k,
+        float* distances,
+        idx_t* labels) const override;
+
+    void reset() override;
+
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    void reconstruct(idx_t key, float* recons) const override;
+
+    size_t remove_ids(const IDSelector& sel) override;
+
+    /* The standalone codec interface */
+    size_t sa_code_size () const override;
+
+    void sa_encode (idx_t n, const float *x,
+                          uint8_t *bytes) const override;
+
+    void sa_decode (idx_t n, const uint8_t *bytes,
+                            float *x) const override;
+
+
+    DistanceComputer * get_distance_computer() const override;
+
+    /******************************************************
+     * Polysemous codes implementation
+     ******************************************************/
+    bool do_polysemous_training; ///< false = standard PQ
+
+    /// parameters used for the polysemous training
+    PolysemousTraining polysemous_training;
+
+    /// how to perform the search in search_core
+    enum Search_type_t {
+        ST_PQ,             ///< asymmetric product quantizer (default)
+        ST_HE,             ///< Hamming distance on codes
+        ST_generalized_HE, ///< nb of same codes
+        ST_SDC,            ///< symmetric product quantizer (SDC)
+        ST_polysemous,     ///< HE filter (using ht) + PQ combination
+        ST_polysemous_generalize,  ///< Filter on generalized Hamming
+    };
+
+    Search_type_t search_type;
+
+    // just encode the sign of the components, instead of using the PQ encoder
+    // used only for the queries
+    bool encode_signs;
+
+    /// Hamming threshold used for polysemy
+    int polysemous_ht;
+
+    // actual polysemous search
+    void search_core_polysemous (idx_t n, const float *x, idx_t k,
+                               float *distances, idx_t *labels) const;
+
+    /// prepare query for a polysemous search, but instead of
+    /// computing the result, just get the histogram of Hamming
+    /// distances. May be computed on a provided dataset if xb != NULL
+    /// @param dist_histogram (M * nbits + 1)
+    void hamming_distance_histogram (idx_t n, const float *x,
+                                     idx_t nb, const float *xb,
+                                     int64_t *dist_histogram);
+
+    /** compute pairwise distances between queries and database
+     *
+     * @param n    nb of query vectors
+     * @param x    query vector, size n * d
+     * @param dis  output distances, size n * ntotal
+     */
+    void hamming_distance_table (idx_t n, const float *x,
+                                 int32_t *dis) const;
+
+};
+
+
+/// statistics are robust to internal threading, but not if
+/// IndexPQ::search is called by multiple threads
+struct IndexPQStats {
+    size_t nq;       // nb of queries run
+    size_t ncode;    // nb of codes visited
+
+    size_t n_hamming_pass; // nb of passed Hamming distance tests (for polysemy)
+
+    IndexPQStats () {reset (); }
+    void reset ();
+};
+
+extern IndexPQStats indexPQ_stats;
+
+
+
+/** Quantizer where centroids are virtual: they are the Cartesian
+ *  product of sub-centroids. */
+struct MultiIndexQuantizer: Index  {
+    ProductQuantizer pq;
+
+    MultiIndexQuantizer (int d,         ///< dimension of the input vectors
+                         size_t M,      ///< number of subquantizers
+                         size_t nbits); ///< number of bit per subvector index
+
+    void train(idx_t n, const float* x) override;
+
+    void search(
+        idx_t n, const float* x, idx_t k,
+        float* distances, idx_t* labels) const override;
+
+    /// add and reset will crash at runtime
+    void add(idx_t n, const float* x) override;
+    void reset() override;
+
+    MultiIndexQuantizer () {}
+
+    void reconstruct(idx_t key, float* recons) const override;
+};
+
+
+/** MultiIndexQuantizer where the PQ assignmnet is performed by sub-indexes
+ */
+struct MultiIndexQuantizer2: MultiIndexQuantizer {
+
+    /// M Indexes on d / M dimensions
+    std::vector<Index*> assign_indexes;
+    bool own_fields;
+
+    MultiIndexQuantizer2 (
+        int d, size_t M, size_t nbits,
+        Index **indexes);
+
+    MultiIndexQuantizer2 (
+        int d, size_t nbits,
+        Index *assign_index_0,
+        Index *assign_index_1);
+
+    void train(idx_t n, const float* x) override;
+
+    void search(
+        idx_t n, const float* x, idx_t k,
+        float* distances, idx_t* labels) const override;
+
+};
+
+
+} // namespace faiss
+
+
+#endif

data/vendor/faiss/IndexPreTransform.cpp

@@ -0,0 +1,288 @@
+/**
+ * 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/IndexPreTransform.h>
+
+#include <cstdio>
+#include <cmath>
+#include <cstring>
+#include <memory>
+
+#include <faiss/utils/utils.h>
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+
+/*********************************************
+ * IndexPreTransform
+ *********************************************/
+
+IndexPreTransform::IndexPreTransform ():
+    index(nullptr), own_fields (false)
+{
+}
+
+
+IndexPreTransform::IndexPreTransform (
+        Index * index):
+    Index (index->d, index->metric_type),
+    index (index), own_fields (false)
+{
+    is_trained = index->is_trained;
+    ntotal = index->ntotal;
+}
+
+
+IndexPreTransform::IndexPreTransform (
+        VectorTransform * ltrans,
+        Index * index):
+    Index (index->d, index->metric_type),
+    index (index), own_fields (false)
+{
+    is_trained = index->is_trained;
+    ntotal = index->ntotal;
+    prepend_transform (ltrans);
+}
+
+void IndexPreTransform::prepend_transform (VectorTransform *ltrans)
+{
+    FAISS_THROW_IF_NOT (ltrans->d_out == d);
+    is_trained = is_trained && ltrans->is_trained;
+    chain.insert (chain.begin(), ltrans);
+    d = ltrans->d_in;
+}
+
+
+IndexPreTransform::~IndexPreTransform ()
+{
+    if (own_fields) {
+        for (int i = 0; i < chain.size(); i++)
+            delete chain[i];
+        delete index;
+    }
+}
+
+
+
+
+void IndexPreTransform::train (idx_t n, const float *x)
+{
+    int last_untrained = 0;
+    if (!index->is_trained) {
+        last_untrained = chain.size();
+    } else {
+        for (int i = chain.size() - 1; i >= 0; i--) {
+            if (!chain[i]->is_trained) {
+                last_untrained = i;
+                break;
+            }
+        }
+    }
+    const float *prev_x = x;
+    ScopeDeleter<float> del;
+
+    if (verbose) {
+        printf("IndexPreTransform::train: training chain 0 to %d\n",
+               last_untrained);
+    }
+
+    for (int i = 0; i <= last_untrained; i++) {
+
+        if (i < chain.size()) {
+            VectorTransform *ltrans = chain [i];
+            if (!ltrans->is_trained) {
+                if (verbose) {
+                    printf("   Training chain component %d/%zd\n",
+                           i, chain.size());
+                    if (OPQMatrix *opqm = dynamic_cast<OPQMatrix*>(ltrans)) {
+                        opqm->verbose = true;
+                    }
+                }
+                ltrans->train (n, prev_x);
+            }
+        } else {
+            if (verbose) {
+                printf("   Training sub-index\n");
+            }
+            index->train (n, prev_x);
+        }
+        if (i == last_untrained) break;
+        if (verbose) {
+            printf("   Applying transform %d/%zd\n",
+                   i, chain.size());
+        }
+
+        float * xt = chain[i]->apply (n, prev_x);
+
+        if (prev_x != x) delete [] prev_x;
+        prev_x = xt;
+        del.set(xt);
+    }
+
+    is_trained = true;
+}
+
+
+const float *IndexPreTransform::apply_chain (idx_t n, const float *x) const
+{
+    const float *prev_x = x;
+    ScopeDeleter<float> del;
+
+    for (int i = 0; i < chain.size(); i++) {
+        float * xt = chain[i]->apply (n, prev_x);
+        ScopeDeleter<float> del2 (xt);
+        del2.swap (del);
+        prev_x = xt;
+    }
+    del.release ();
+    return prev_x;
+}
+
+void IndexPreTransform::reverse_chain (idx_t n, const float* xt, float* x) const
+{
+    const float* next_x = xt;
+    ScopeDeleter<float> del;
+
+    for (int i = chain.size() - 1; i >= 0; i--) {
+        float* prev_x = (i == 0) ? x : new float [n * chain[i]->d_in];
+        ScopeDeleter<float> del2 ((prev_x == x) ? nullptr : prev_x);
+        chain [i]->reverse_transform (n, next_x, prev_x);
+        del2.swap (del);
+        next_x = prev_x;
+    }
+}
+
+void IndexPreTransform::add (idx_t n, const float *x)
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    const float *xt = apply_chain (n, x);
+    ScopeDeleter<float> del(xt == x ? nullptr : xt);
+    index->add (n, xt);
+    ntotal = index->ntotal;
+}
+
+void IndexPreTransform::add_with_ids (idx_t n, const float * x,
+                                      const idx_t *xids)
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    const float *xt = apply_chain (n, x);
+    ScopeDeleter<float> del(xt == x ? nullptr : xt);
+    index->add_with_ids (n, xt, xids);
+    ntotal = index->ntotal;
+}
+
+
+
+
+void IndexPreTransform::search (idx_t n, const float *x, idx_t k,
+                               float *distances, idx_t *labels) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    const float *xt = apply_chain (n, x);
+    ScopeDeleter<float> del(xt == x ? nullptr : xt);
+    index->search (n, xt, k, distances, labels);
+}
+
+void IndexPreTransform::range_search (idx_t n, const float* x, float radius,
+                                      RangeSearchResult* result) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+    const float *xt = apply_chain (n, x);
+    ScopeDeleter<float> del(xt == x ? nullptr : xt);
+    index->range_search (n, xt, radius, result);
+}
+
+
+
+void IndexPreTransform::reset () {
+    index->reset();
+    ntotal = 0;
+}
+
+size_t IndexPreTransform::remove_ids (const IDSelector & sel) {
+    size_t nremove = index->remove_ids (sel);
+    ntotal = index->ntotal;
+    return nremove;
+}
+
+
+void IndexPreTransform::reconstruct (idx_t key, float * recons) const
+{
+    float *x = chain.empty() ? recons : new float [index->d];
+    ScopeDeleter<float> del (recons == x ? nullptr : x);
+    // Initial reconstruction
+    index->reconstruct (key, x);
+
+    // Revert transformations from last to first
+    reverse_chain (1, x, recons);
+}
+
+
+void IndexPreTransform::reconstruct_n (idx_t i0, idx_t ni, float *recons) const
+{
+    float *x = chain.empty() ? recons : new float [ni * index->d];
+    ScopeDeleter<float> del (recons == x ? nullptr : x);
+    // Initial reconstruction
+    index->reconstruct_n (i0, ni, x);
+
+    // Revert transformations from last to first
+    reverse_chain (ni, x, recons);
+}
+
+
+void IndexPreTransform::search_and_reconstruct (
+      idx_t n, const float *x, idx_t k,
+      float *distances, idx_t *labels, float* recons) const
+{
+    FAISS_THROW_IF_NOT (is_trained);
+
+    const float* xt = apply_chain (n, x);
+    ScopeDeleter<float> del ((xt == x) ? nullptr : xt);
+
+    float* recons_temp = chain.empty() ? recons : new float [n * k * index->d];
+    ScopeDeleter<float> del2 ((recons_temp == recons) ? nullptr : recons_temp);
+    index->search_and_reconstruct (n, xt, k, distances, labels, recons_temp);
+
+    // Revert transformations from last to first
+    reverse_chain (n * k, recons_temp, recons);
+}
+
+size_t IndexPreTransform::sa_code_size () const
+{
+    return index->sa_code_size ();
+}
+
+void IndexPreTransform::sa_encode (idx_t n, const float *x,
+                                         uint8_t *bytes) const
+{
+    if (chain.empty()) {
+        index->sa_encode (n, x, bytes);
+    } else {
+        const float *xt = apply_chain (n, x);
+        ScopeDeleter<float> del(xt == x ? nullptr : xt);
+        index->sa_encode (n, xt, bytes);
+    }
+}
+
+void IndexPreTransform::sa_decode (idx_t n, const uint8_t *bytes,
+                                           float *x) const
+{
+    if (chain.empty()) {
+        index->sa_decode (n, bytes, x);
+    } else {
+        std::unique_ptr<float []> x1 (new float [index->d * n]);
+        index->sa_decode (n, bytes, x1.get());
+        // Revert transformations from last to first
+        reverse_chain (n, x1.get(), x);
+    }
+}
+
+
+
+} // namespace faiss