Add TurboQuant index and benchmark support

Author: Mistobaan

.gitignore

@@ -31,3 +31,5 @@ faiss/python/swigfaiss_sve.swig
 # Python package build outputs
 /dist/
 /*.egg-info/
+/build-conda/
+/benchs/sift1M/

benchs/__init__.py

@@ -10,29 +10,80 @@
 
 try:
     from faiss.contrib.datasets_fb import \
-        DatasetSIFT1M, DatasetDeep1B, DatasetBigANN
+        DatasetSIFT1M, DatasetDeep1B, DatasetBigANN, DatasetGlove
 except ImportError:
     from faiss.contrib.datasets import \
-        DatasetSIFT1M, DatasetDeep1B, DatasetBigANN
+        DatasetSIFT1M, DatasetDeep1B, DatasetBigANN, DatasetGlove, DatasetDBpedia1536_1M, DatasetDBpedia3072_1M
 
 
-def eval_codec(q, xq, xb, gt):
+TURBOQUANT_OPTIONS = {"turboquant", "tq"}
+RABITQ_OPTIONS = {"rabitq", "rbq"}
+
+
+def get_metric_type(ds):
+    if ds.metric == "IP":
+        return faiss.METRIC_INNER_PRODUCT
+    if ds.metric == "L2":
+        return faiss.METRIC_L2
+    raise RuntimeError(f"unsupported dataset metric {ds.metric}")
+
+
+def get_training_vectors(ds, xb, maxtrain):
+    try:
+        return ds.get_train(maxtrain=maxtrain)
+    except NotImplementedError:
+        print("No training set: training on database")
+        return xb[:maxtrain]
+
+
+def encode(codec, x):
+    if hasattr(codec, "compute_codes") and hasattr(codec, "decode"):
+        return codec.compute_codes(x)
+    if hasattr(codec, "sa_encode") and hasattr(codec, "sa_decode"):
+        return codec.sa_encode(x)
+    raise TypeError(f"unsupported codec type {type(codec).__name__}")
+
+
+def decode(codec, codes):
+    if hasattr(codec, "compute_codes") and hasattr(codec, "decode"):
+        return codec.decode(codes)
+    if hasattr(codec, "sa_encode") and hasattr(codec, "sa_decode"):
+        return codec.sa_decode(codes)
+    raise TypeError(f"unsupported codec type {type(codec).__name__}")
+
+
+def get_code_size(codec):
+    if hasattr(codec, "code_size"):
+        return int(codec.code_size)
+    if hasattr(codec, "sa_code_size"):
+        return int(codec.sa_code_size())
+    return None
+
+
+def eval_codec(q, xq, xb, gt, metric_type):
     t0 = time.time()
-    codes = q.compute_codes(xb)
+    codes = encode(q, xb)
     t1 = time.time()
-    xb_decoded = q.decode(codes)
+    xb_decoded = decode(q, codes)
     recons_err = ((xb - xb_decoded) ** 2).sum() / xb.shape[0]
     # for compatibility with the codec benchmarks
     err_compat = np.linalg.norm(xb - xb_decoded, axis=1).mean()
-    xq_decoded = q.decode(q.compute_codes(xq))
-    D, I = faiss.knn(xq_decoded, xb_decoded, 1)
+    xq_decoded = decode(q, encode(q, xq))
+    D, I = faiss.knn(xq_decoded, xb_decoded, 1, metric=metric_type)
     recall = (I[:, 0] == gt[:, 0]).sum() / nq
+    code_size = get_code_size(q)
+    code_size_s = (
+        f" code_size: {code_size} B/vector"
+        if code_size is not None
+        else ""
+    )
     print(
         f"\tencode time: {t1 - t0:.3f} reconstruction error: {recons_err:.3f} "
-        f"1-recall@1: {recall:.4f} recons_err_compat {err_compat:.3f}")
+        f"recall@1: {recall:.4f} recons_err_compat {err_compat:.3f}"
+        f"{code_size_s}")
 
 
-def eval_quantizer(q, xq, xb, gt, xt, variants=None):
+def eval_quantizer(q, xq, xb, gt, xt, metric_type, variants=None):
     if variants is None:
         variants = [(None, None)]
     t0 = time.time()
@@ -53,20 +104,41 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
                 getattr(q, name)  # make sure field exists
                 setattr(q, name, val)
 
-        eval_codec(q, xq, xb, gt)
+        eval_codec(q, xq, xb, gt, metric_type)
 
 
 todo = sys.argv[1:]
 
-if len(todo) > 0 and "deep1M" in todo[0]:
-    ds = DatasetDeep1B(10**6)
+dataset_name = "sift1M"
+if len(todo) > 0 and todo[0] in (
+    "sift1M",
+    "deep1M",
+    "bigann1M",
+    "glove",
+    "dbpedia-1536-1M",
+    "dbpedia-3072-1M",
+):
+    dataset_name = todo[0]
     del todo[0]
-elif len(todo) > 0 and "bigann1M" in todo[0]:
+
+if dataset_name == "deep1M":
+    ds = DatasetDeep1B(10**6)
+elif dataset_name == "bigann1M":
     ds = DatasetBigANN(nb_M=1)
-    del todo[0]
+elif dataset_name == "glove":
+    ds = DatasetGlove()
+elif dataset_name == "dbpedia-1536-1M":
+    ds = DatasetDBpedia1536_1M()
+elif dataset_name == "dbpedia-3072-1M":
+    ds = DatasetDBpedia3072_1M()
 else:
     ds = DatasetSIFT1M()
 
+M = None
+nsplits = None
+Msub = None
+nbits = None
+
 if len(todo) > 0:
     if todo[0].count("x") == 1:
         M, nbits = [int(x) for x in todo[0].split("x")]
@@ -75,15 +147,42 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
         nsplits, Msub, nbits = [int(x) for x in todo[0].split("x")]
         M = nsplits * Msub
         del todo[0]
+    elif todo[0].isdigit():
+        nbits = int(todo[0])
+        del todo[0]
+
+selected_options = set(todo)
+
+if nbits is None:
+    raise RuntimeError(
+        "expected a codec bit specification: Mxnbits, nsplitsxMsubxnbits, "
+        "or plain nbits for turboquant/rabitq"
+    )
+
+if selected_options & TURBOQUANT_OPTIONS and not 1 <= nbits <= 8:
+    raise RuntimeError("TurboQuant supports nbits in [1, 8]")
+
+if selected_options & RABITQ_OPTIONS and not 1 <= nbits <= 9:
+    raise RuntimeError("RaBitQ supports nbits in [1, 9]")
+
+if M is None and selected_options & {"pq", "opq", "rq", "rq_lut", "lsq", "lsq-gpu"}:
+    raise RuntimeError("expected Mxnbits for pq/opq/rq/rq_lut/lsq benchmarks")
+
+if M is None and selected_options & {"prq", "plsq"}:
+    raise RuntimeError("expected nsplitsxMsubxnbits for prq/plsq benchmarks")
 
 maxtrain = max(100 << nbits, 10**5)
-print(f"eval on {M}x{nbits} maxtrain={maxtrain}")
+if M is None:
+    print(f"eval on {dataset_name} {nbits}-bit maxtrain={maxtrain}")
+else:
+    print(f"eval on {dataset_name} {M}x{nbits} maxtrain={maxtrain}")
 
 xq = ds.get_queries()
 xb = ds.get_database()
 gt = ds.get_groundtruth()
+metric_type = get_metric_type(ds)
 
-xt = ds.get_train(maxtrain=maxtrain)
+xt = get_training_vectors(ds, xb, maxtrain=maxtrain)
 
 nb, d = xb.shape
 nq, d = xq.shape
@@ -97,12 +196,12 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
     ngpus = faiss.get_num_gpus()
     lsq.icm_encoder_factory = faiss.GpuIcmEncoderFactory(ngpus)
     lsq.verbose = True
-    eval_quantizer(lsq, xb, xt, 'lsq-gpu')
+    eval_quantizer(lsq, xq, xb, gt, xt, metric_type)
 
 if 'pq' in todo:
     pq = faiss.ProductQuantizer(d, M, nbits)
     print("===== PQ")
-    eval_quantizer(pq, xq, xb, gt, xt)
+    eval_quantizer(pq, xq, xb, gt, xt, metric_type)
 
 if 'opq' in todo:
     d2 = ((d + M - 1) // M) * M
@@ -114,19 +213,19 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
     xt2 = opq.apply(xt)
     pq = faiss.ProductQuantizer(d2, M, nbits)
     print("===== PQ")
-    eval_quantizer(pq, xq2, xb2, gt, xt2)
+    eval_quantizer(pq, xq2, xb2, gt, xt2, metric_type)
 
 if 'prq' in todo:
     print(f"===== PRQ{nsplits}x{Msub}x{nbits}")
     prq = faiss.ProductResidualQuantizer(d, nsplits, Msub, nbits)
     variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32)]
-    eval_quantizer(prq, xq, xb, gt, xt, variants=variants)
+    eval_quantizer(prq, xq, xb, gt, xt, metric_type, variants=variants)
 
 if 'plsq' in todo:
     print(f"===== PLSQ{nsplits}x{Msub}x{nbits}")
     plsq = faiss.ProductLocalSearchQuantizer(d, nsplits, Msub, nbits)
     variants = [("encode_ils_iters", i) for i in (2, 3, 4, 8, 16)]
-    eval_quantizer(plsq, xq, xb, gt, xt, variants=variants)
+    eval_quantizer(plsq, xq, xb, gt, xt, metric_type, variants=variants)
 
 if 'rq' in todo:
     print("===== RQ")
@@ -136,7 +235,7 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
     # rq.train_type = faiss.ResidualQuantizer.Train_default
     # rq.verbose = True
     variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32)]
-    eval_quantizer(rq, xq, xb, gt, xt, variants=variants)
+    eval_quantizer(rq, xq, xb, gt, xt, metric_type, variants=variants)
 
 if 'rq_lut' in todo:
     print("===== RQ")
@@ -148,10 +247,34 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
     # rq.train_type = faiss.ResidualQuantizer.Train_default
     # rq.verbose = True
     variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32, 64)]
-    eval_quantizer(rq, xq, xb, gt, xt, variants=variants)
+    eval_quantizer(rq, xq, xb, gt, xt, metric_type, variants=variants)
 
 if 'lsq' in todo:
     print("===== LSQ")
     lsq = faiss.LocalSearchQuantizer(d, M, nbits)
     variants = [("encode_ils_iters", i) for i in (2, 3, 4, 8, 16)]
-    eval_quantizer(lsq, xq, xb, gt, xt, variants=variants)
+    eval_quantizer(lsq, xq, xb, gt, xt, metric_type, variants=variants)
+
+if selected_options & TURBOQUANT_OPTIONS:
+    print("===== TurboQuant")
+    store_norm = dataset_name != "glove"
+    if hasattr(faiss, "IndexTurboQuantMSE"):
+        tq = faiss.IndexTurboQuantMSE(d, nbits, metric_type, 12345, store_norm)
+    elif hasattr(faiss, "TurboQuantizer"):
+        tq = faiss.TurboQuantizer(d, nbits, 12345, store_norm)
+    else:
+        raise RuntimeError(
+            "TurboQuant is not available in this faiss Python build. "
+            "Rebuild the Python bindings so TurboQuant symbols are exported."
+        )
+    eval_quantizer(tq, xq, xb, gt, xt, metric_type)
+
+if selected_options & RABITQ_OPTIONS:
+    print("===== RaBitQ")
+    if not hasattr(faiss, "IndexRaBitQ"):
+        raise RuntimeError(
+            "RaBitQ is not available in this faiss Python build. "
+            "Rebuild the Python bindings so RaBitQ symbols are exported."
+        )
+    rbq = faiss.IndexRaBitQ(d, metric_type, nbits)
+    eval_quantizer(rbq, xq, xb, gt, xt, metric_type)