add reconstruct support to additive quantizers (facebookresearch#3752)

Summary: Pull Request resolved: facebookresearch#3752 add reconstruct support to additive quantizers. fixes github issue: 2422 Reviewed By: asadoughi Differential Revision: D61255049 fbshipit-source-id: 09a0edae7fc24295a686d332e2d052e37372d2c0
aalekhpatel07 · Oct 17, 2024 · 70da3c3 · 70da3c3
1 parent 3b34f59
commit 70da3c3
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 0 deletions.
diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/faiss/IndexIVFAdditiveQuantizer.cpp
@@ -116,6 +116,21 @@ void IndexIVFAdditiveQuantizer::sa_decode(
     }
 }
 
+void IndexIVFAdditiveQuantizer::reconstruct_from_offset(
+        int64_t list_no,
+        int64_t offset,
+        float* recons) const {
+    const uint8_t* code = invlists->get_single_code(list_no, offset);
+    aq->decode(code, recons, 1);
+    if (by_residual) {
+        std::vector<float> centroid(d);
+        quantizer->reconstruct(list_no, centroid.data());
+        for (int i = 0; i < d; ++i) {
+            recons[i] += centroid[i];
+        }
+    }
+}
+
 IndexIVFAdditiveQuantizer::~IndexIVFAdditiveQuantizer() = default;
 
 /*********************************************

diff --git a/faiss/IndexIVFAdditiveQuantizer.h b/faiss/IndexIVFAdditiveQuantizer.h
@@ -56,6 +56,9 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
 
     void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
 
+    void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
+            const override;
+
     ~IndexIVFAdditiveQuantizer() override;
 };
 

diff --git a/tests/test_index.py b/tests/test_index.py
@@ -493,6 +493,23 @@ def test_IndexIVFPQ(self):
 
         self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
 
+    def test_IndexIVFRQ(self):
+        d = 32
+        nb = 1000
+        nt = 1500
+        nq = 200
+
+        (xt, xb, xq) = get_dataset(d, nb, nt, nq)
+
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFResidualQuantizer(quantizer, d, 32, 8, 8)
+        index.cp.min_points_per_centroid = 5    # quiet warning
+        index.nprobe = 4
+        index.train(xt)
+        index.add(xb)
+
+        self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
+
     def test_MultiIndex(self):
         d = 32
         nb = 1000