faiss 0.3.0 → 0.3.2

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp

@@ -125,51 +125,27 @@ IndexIVFAdditiveQuantizerFastScan::IndexIVFAdditiveQuantizerFastScan() {
     is_trained = false;
 }
 
-IndexIVFAdditiveQuantizerFastScan::~IndexIVFAdditiveQuantizerFastScan() {}
+IndexIVFAdditiveQuantizerFastScan::~IndexIVFAdditiveQuantizerFastScan() =
+        default;
 
 /*********************************************************
  * Training
  *********************************************************/
 
-void IndexIVFAdditiveQuantizerFastScan::train_residual(
+idx_t IndexIVFAdditiveQuantizerFastScan::train_encoder_num_vectors() const {
+    return max_train_points;
+}
+
+void IndexIVFAdditiveQuantizerFastScan::train_encoder(
         idx_t n,
-        const float* x_in) {
+        const float* x,
+        const idx_t* assign) {
     if (aq->is_trained) {
         return;
     }
 
-    const int seed = 0x12345;
-    size_t nt = n;
-    const float* x = fvecs_maybe_subsample(
-            d, &nt, max_train_points, x_in, verbose, seed);
-    n = nt;
     if (verbose) {
-        printf("training additive quantizer on %zd vectors\n", nt);
-    }
-    aq->verbose = verbose;
-
-    std::unique_ptr<float[]> del_x;
-    if (x != x_in) {
-        del_x.reset((float*)x);
-    }
-
-    const float* trainset;
-    std::vector<float> residuals(n * d);
-    std::vector<idx_t> assign(n);
-
-    if (by_residual) {
-        if (verbose) {
-            printf("computing residuals\n");
-        }
-        quantizer->assign(n, x, assign.data());
-        residuals.resize(n * d);
-        for (idx_t i = 0; i < n; i++) {
-            quantizer->compute_residual(
-                    x + i * d, residuals.data() + i * d, assign[i]);
-        }
-        trainset = residuals.data();
-    } else {
-        trainset = x;
+        printf("training additive quantizer on %d vectors\n", int(n));
     }
 
     if (verbose) {
@@ -181,17 +157,16 @@ void IndexIVFAdditiveQuantizerFastScan::train_residual(
                d);
     }
     aq->verbose = verbose;
-    aq->train(n, trainset);
+    aq->train(n, x);
 
     // train norm quantizer
     if (by_residual && metric_type == METRIC_L2) {
         std::vector<float> decoded_x(n * d);
         std::vector<uint8_t> x_codes(n * aq->code_size);
-        aq->compute_codes(residuals.data(), x_codes.data(), n);
+        aq->compute_codes(x, x_codes.data(), n);
         aq->decode(x_codes.data(), decoded_x.data(), n);
 
         // add coarse centroids
-        FAISS_THROW_IF_NOT(assign.size() == n);
         std::vector<float> centroid(d);
         for (idx_t i = 0; i < n; i++) {
             auto xi = decoded_x.data() + i * d;
@@ -236,7 +211,8 @@ void IndexIVFAdditiveQuantizerFastScan::estimate_norm_scale(
 
     size_t index_nprobe = nprobe;
     nprobe = 1;
-    compute_LUT(n, x, coarse_ids.data(), coarse_dis.data(), dis_tables, biases);
+    CoarseQuantized cq{index_nprobe, coarse_dis.data(), coarse_ids.data()};
+    compute_LUT(n, x, cq, dis_tables, biases);
     nprobe = index_nprobe;
 
     float scale = 0;
@@ -338,11 +314,8 @@ void IndexIVFAdditiveQuantizerFastScan::search(
     }
 
     NormTableScaler scaler(norm_scale);
-    if (metric_type == METRIC_L2) {
-        search_dispatch_implem<true>(n, x, k, distances, labels, scaler);
-    } else {
-        search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
-    }
+    IndexIVFFastScan::CoarseQuantized cq{nprobe};
+    search_dispatch_implem(n, x, k, distances, labels, cq, &scaler);
 }
 
 /*********************************************************
@@ -408,12 +381,12 @@ bool IndexIVFAdditiveQuantizerFastScan::lookup_table_is_3d() const {
 void IndexIVFAdditiveQuantizerFastScan::compute_LUT(
         size_t n,
         const float* x,
-        const idx_t* coarse_ids,
-        const float*,
+        const CoarseQuantized& cq,
         AlignedTable<float>& dis_tables,
         AlignedTable<float>& biases) const {
     const size_t dim12 = ksub * M;
     const size_t ip_dim12 = aq->M * ksub;
+    const size_t nprobe = cq.nprobe;
 
     dis_tables.resize(n * dim12);
 
@@ -434,7 +407,7 @@ void IndexIVFAdditiveQuantizerFastScan::compute_LUT(
 #pragma omp for
             for (idx_t ij = 0; ij < n * nprobe; ij++) {
                 int i = ij / nprobe;
-                quantizer->reconstruct(coarse_ids[ij], c);
+                quantizer->reconstruct(cq.ids[ij], c);
                 biases[ij] = coef * fvec_inner_product(c, x + i * d, d);
             }
         }

data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h

@@ -63,7 +63,9 @@ struct IndexIVFAdditiveQuantizerFastScan : IndexIVFFastScan {
             const IndexIVFAdditiveQuantizer& orig,
             int bbs = 32);
 
-    void train_residual(idx_t n, const float* x) override;
+    void train_encoder(idx_t n, const float* x, const idx_t* assign) override;
+
+    idx_t train_encoder_num_vectors() const override;
 
     void estimate_norm_scale(idx_t n, const float* x);
 
@@ -91,8 +93,7 @@ struct IndexIVFAdditiveQuantizerFastScan : IndexIVFFastScan {
     void compute_LUT(
             size_t n,
             const float* x,
-            const idx_t* coarse_ids,
-            const float* coarse_dis,
+            const CoarseQuantized& cq,
             AlignedTable<float>& dis_tables,
             AlignedTable<float>& biases) const override;