faiss/benchs/bench_quantizer.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import sys
import faiss
import time
import numpy as np

try:
    from faiss.contrib.datasets_fb import (
        DatasetSIFT1M,
        DatasetDeep1B,
        DatasetBigANN,
    )
except ImportError:
    from faiss.contrib.datasets import (
        DatasetSIFT1M,
        DatasetDeep1B,
        DatasetBigANN,
    )


def eval_codec(q, xq, xb, gt):
    t0 = time.time()
    codes = q.compute_codes(xb)
    t1 = time.time()
    xb_decoded = q.decode(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)
    recall = (I[:, 0] == gt[:, 0]).sum() / nq
    print(
        f"\tencode time: {t1 - t0:.3f} reconstruction error: {recons_err:.3f} "
        f"1-recall@1: {recall:.4f} recons_err_compat {err_compat:.3f}"
    )


def eval_quantizer(q, xq, xb, gt, xt, variants=None):
    if variants is None:
        variants = [(None, None)]
    t0 = time.time()
    q.train(xt)
    t1 = time.time()
    train_t = t1 - t0
    print(f"\ttraining time: {train_t:.3f} s")
    for name, val in variants:
        if name is not None:
            print(f"{name}={val}")

            if isinstance(q, faiss.ProductAdditiveQuantizer):
                for i in range(q.nsplits):
                    subq = faiss.downcast_Quantizer(q.subquantizer(i))
                    getattr(subq, name)
                    setattr(subq, name, val)
            else:
                getattr(q, name)  # make sure field exists
                setattr(q, name, val)

        eval_codec(q, xq, xb, gt)


todo = sys.argv[1:]

if len(todo) > 0 and "deep1M" in todo[0]:
    ds = DatasetDeep1B(10**6)
    del todo[0]
elif len(todo) > 0 and "bigann1M" in todo[0]:
    ds = DatasetBigANN(nb_M=1)
    del todo[0]
else:
    ds = DatasetSIFT1M()

if len(todo) > 0:
    if todo[0].count("x") == 1:
        M, nbits = [int(x) for x in todo[0].split("x")]
        del todo[0]
    elif todo[0].count("x") == 2:
        nsplits, Msub, nbits = [int(x) for x in todo[0].split("x")]
        M = nsplits * Msub
        del todo[0]

maxtrain = max(100 << nbits, 10**5)
print(f"eval on {M}x{nbits} maxtrain={maxtrain}")

xq = ds.get_queries()
xb = ds.get_database()
gt = ds.get_groundtruth()

xt = ds.get_train(maxtrain=maxtrain)

nb, d = xb.shape
nq, d = xq.shape
nt, d = xt.shape


# fastest to slowest

if "lsq-gpu" in todo:
    lsq = faiss.LocalSearchQuantizer(d, M, nbits)
    ngpus = faiss.get_num_gpus()
    lsq.icm_encoder_factory = faiss.GpuIcmEncoderFactory(ngpus)
    lsq.verbose = True
    eval_quantizer(lsq, xb, xt, "lsq-gpu")

if "pq" in todo:
    pq = faiss.ProductQuantizer(d, M, nbits)
    print("===== PQ")
    eval_quantizer(pq, xq, xb, gt, xt)

if "opq" in todo:
    d2 = ((d + M - 1) // M) * M
    print("OPQ d2=", d2)
    opq = faiss.OPQMatrix(d, M, d2)
    opq.train(xt)
    xq2 = opq.apply(xq)
    xb2 = opq.apply(xb)
    xt2 = opq.apply(xt)
    pq = faiss.ProductQuantizer(d2, M, nbits)
    print("===== PQ")
    eval_quantizer(pq, xq2, xb2, gt, xt2)

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)

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)

if "rq" in todo:
    print("===== RQ")
    rq = faiss.ResidualQuantizer(
        d,
        M,
        nbits,
    )
    rq.max_beam_size
    rq.max_beam_size = 30  # for compatibility with older runs
    # 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)

if "rq_lut" in todo:
    print("===== RQ")
    rq = faiss.ResidualQuantizer(
        d,
        M,
        nbits,
    )
    rq.max_beam_size
    rq.max_beam_size = 30  # for compatibility with older runs
    rq.use_beam_LUT
    rq.use_beam_LUT = 1
    # 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)

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)