faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/ResidualQuantizer.h

@@ -24,22 +24,38 @@ namespace faiss {
 
 struct ResidualQuantizer : AdditiveQuantizer {
     /// initialization
-    enum train_type_t {
-        Train_default,         ///< regular k-means
-        Train_progressive_dim, ///< progressive dim clustering
-    };
 
-    // set this bit on train_type if beam is to be trained only on the
-    // first element of the beam (faster but less accurate)
-    static const int Train_top_beam = 1024;
+    //  Was enum but that does not work so well with bitmasks
+    using train_type_t = int;
+
+    /// Binary or of the Train_* flags below
     train_type_t train_type;
 
+    /// regular k-means (minimal amount of computation)
+    static const int Train_default = 0;
+
+    /// progressive dim clustering (set by default)
+    static const int Train_progressive_dim = 1;
+
+    /// do a few iterations of codebook refinement after first level estimation
+    static const int Train_refine_codebook = 2;
+
+    /// number of iterations for codebook refinement.
+    int niter_codebook_refine;
+
+    /** set this bit on train_type if beam is to be trained only on the
+     *  first element of the beam (faster but less accurate) */
+    static const int Train_top_beam = 1024;
+
+    /** set this bit to *not* autmatically compute the codebook tables
+     * after training */
+    static const int Skip_codebook_tables = 2048;
+
     /// beam size used for training and for encoding
     int max_beam_size;
 
-    /// distance matrixes with beam search can get large, so use this
-    /// to batch computations at encoding time.
-    size_t max_mem_distances;
+    /// use LUT for beam search
+    int use_beam_LUT;
 
     /// clustering parameters
     ProgressiveDimClusteringParameters cp;
@@ -47,24 +63,46 @@ struct ResidualQuantizer : AdditiveQuantizer {
     /// if non-NULL, use this index for assignment
     ProgressiveDimIndexFactory* assign_index_factory;
 
-    ResidualQuantizer(size_t d, const std::vector<size_t>& nbits);
+    ResidualQuantizer(
+            size_t d,
+            const std::vector<size_t>& nbits,
+            Search_type_t search_type = ST_decompress);
 
     ResidualQuantizer(
-            size_t d,      /* dimensionality of the input vectors */
-            size_t M,      /* number of subquantizers */
-            size_t nbits); /* number of bit per subvector index */
+            size_t d,     /* dimensionality of the input vectors */
+            size_t M,     /* number of subquantizers */
+            size_t nbits, /* number of bit per subvector index */
+            Search_type_t search_type = ST_decompress);
 
     ResidualQuantizer();
 
-    // Train the residual quantizer
+    /// Train the residual quantizer
     void train(size_t n, const float* x) override;
 
+    /// Copy the M codebook levels from other, starting from skip_M
+    void initialize_from(const ResidualQuantizer& other, int skip_M = 0);
+
+    /** Encode the vectors and compute codebook that minimizes the quantization
+     * error on these codes
+     *
+     * @param x      training vectors, size n * d
+     * @param n      nb of training vectors, n >= total_codebook_size
+     * @return       returns quantization error for the new codebook with old
+     * codes
+     */
+    float retrain_AQ_codebook(size_t n, const float* x);
+
     /** Encode a set of vectors
      *
      * @param x      vectors to encode, size n * d
      * @param codes  output codes, size n * code_size
+     * @param centroids  centroids to be added to x, size n * d
      */
-    void compute_codes(const float* x, uint8_t* codes, size_t n) const override;
+    void compute_codes_add_centroids(
+            const float* x,
+            uint8_t* codes,
+            size_t n,
+            const float* centroids = nullptr) const override;
 
     /** lower-level encode function
      *
@@ -85,12 +123,32 @@ struct ResidualQuantizer : AdditiveQuantizer {
             float* new_residuals = nullptr,
             float* new_distances = nullptr) const;
 
+    void refine_beam_LUT(
+            size_t n,
+            const float* query_norms,
+            const float* query_cp,
+            int new_beam_size,
+            int32_t* new_codes,
+            float* new_distances = nullptr) const;
+
     /** Beam search can consume a lot of memory. This function estimates the
      * amount of mem used by refine_beam to adjust the batch size
      *
      * @param beam_size  if != -1, override the beam size
      */
     size_t memory_per_point(int beam_size = -1) const;
+
+    /** Cross products used in codebook tables
+     *
+     * These are used to keep trak of norms of centroids.
+     */
+    void compute_codebook_tables();
+
+    /// dot products of all codebook vectors with each other
+    /// size total_codebook_size * total_codebook_size
+    std::vector<float> codebook_cross_products;
+    /// norms of all vectors
+    std::vector<float> cent_norms;
 };
 
 /** Encode a residual by sampling from a centroid table.
@@ -127,4 +185,24 @@ void beam_search_encode_step(
         float* new_distances,
         Index* assign_index = nullptr);
 
+/** Encode a set of vectors using their dot products with the codebooks
+ *
+ */
+void beam_search_encode_step_tab(
+        size_t K,
+        size_t n,
+        size_t beam_size,                  // input sizes
+        const float* codebook_cross_norms, // size K * ldc
+        size_t ldc,                        // >= K
+        const uint64_t* codebook_offsets,  // m
+        const float* query_cp,             // size n * ldqc
+        size_t ldqc,                       // >= K
+        const float* cent_norms_i,         // size K
+        size_t m,
+        const int32_t* codes,   // n * beam_size * m
+        const float* distances, // n * beam_size
+        size_t new_beam_size,
+        int32_t* new_codes,    // n * new_beam_size * (m + 1)
+        float* new_distances); // n * new_beam_size
+
 }; // namespace faiss

data/vendor/faiss/faiss/impl/ResultHandler.h

@@ -413,4 +413,100 @@ struct RangeSearchResultHandler {
     }
 };
 
+/*****************************************************************
+ * Single best result handler.
+ * Tracks the only best result, thus avoiding storing
+ * some temporary data in memory.
+ *****************************************************************/
+
+template <class C>
+struct SingleBestResultHandler {
+    using T = typename C::T;
+    using TI = typename C::TI;
+
+    int nq;
+    // contains exactly nq elements
+    T* dis_tab;
+    // contains exactly nq elements
+    TI* ids_tab;
+
+    SingleBestResultHandler(size_t nq, T* dis_tab, TI* ids_tab)
+            : nq(nq), dis_tab(dis_tab), ids_tab(ids_tab) {}
+
+    struct SingleResultHandler {
+        SingleBestResultHandler& hr;
+
+        T min_dis;
+        TI min_idx;
+        size_t current_idx = 0;
+
+        SingleResultHandler(SingleBestResultHandler& hr) : hr(hr) {}
+
+        /// begin results for query # i
+        void begin(const size_t current_idx) {
+            this->current_idx = current_idx;
+            min_dis = HUGE_VALF;
+            min_idx = 0;
+        }
+
+        /// add one result for query i
+        void add_result(T dis, TI idx) {
+            if (C::cmp(min_dis, dis)) {
+                min_dis = dis;
+                min_idx = idx;
+            }
+        }
+
+        /// series of results for query i is done
+        void end() {
+            hr.dis_tab[current_idx] = min_dis;
+            hr.ids_tab[current_idx] = min_idx;
+        }
+    };
+
+    size_t i0, i1;
+
+    /// begin
+    void begin_multiple(size_t i0, size_t i1) {
+        this->i0 = i0;
+        this->i1 = i1;
+
+        for (size_t i = i0; i < i1; i++) {
+            this->dis_tab[i] = HUGE_VALF;
+        }
+    }
+
+    /// add results for query i0..i1 and j0..j1
+    void add_results(size_t j0, size_t j1, const T* dis_tab) {
+        for (int64_t i = i0; i < i1; i++) {
+            const T* dis_tab_i = dis_tab + (j1 - j0) * (i - i0) - j0;
+
+            auto& min_distance = this->dis_tab[i];
+            auto& min_index = this->ids_tab[i];
+
+            for (size_t j = j0; j < j1; j++) {
+                const T distance = dis_tab_i[j];
+
+                if (C::cmp(min_distance, distance)) {
+                    min_distance = distance;
+                    min_index = j;
+                }
+            }
+        }
+    }
+
+    void add_result(const size_t i, const T dis, const TI idx) {
+        auto& min_distance = this->dis_tab[i];
+        auto& min_index = this->ids_tab[i];
+
+        if (C::cmp(min_distance, dis)) {
+            min_distance = dis;
+            min_index = idx;
+        }
+    }
+
+    /// series of results for queries i0..i1 is done
+    void end_multiple() {}
+};
+
 } // namespace faiss