faiss/gpu/test/test_gpu_index_ivfsq.py

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

#! /usr/bin/env python3

from __future__ import print_function
import unittest
import numpy as np
import faiss

def make_t(num, d, clamp=False):
    rs = np.random.RandomState(123)
    x = rs.rand(num, d).astype('float32')
    if clamp:
        x = (x * 255).astype('uint8').astype('float32')
    return x

def make_indices_copy_from_cpu(nlist, d, qtype, by_residual, metric, clamp):
    to_train = make_t(10000, d, clamp)

    quantizer_cp = faiss.IndexFlat(d, metric)
    idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
                                            qtype, metric, by_residual)

    idx_cpu.train(to_train)
    idx_cpu.add(to_train)

    res = faiss.StandardGpuResources()
    res.noTempMemory()
    idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, idx_cpu)

    return idx_cpu, idx_gpu


def make_indices_copy_from_gpu(nlist, d, qtype, by_residual, metric, clamp):
    to_train = make_t(10000, d, clamp)

    res = faiss.StandardGpuResources()
    res.noTempMemory()
    idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist,
                                               qtype, metric, by_residual)
    idx_gpu.train(to_train)
    idx_gpu.add(to_train)

    quantizer_cp = faiss.IndexFlat(d, metric)
    idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
                                            qtype, metric, by_residual)
    idx_gpu.copyTo(idx_cpu)

    return idx_cpu, idx_gpu


def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
    to_train = make_t(10000, d, clamp)

    quantizer_cp = faiss.IndexFlat(d, metric)
    idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
                                            qtype, metric, by_residual)
    assert(by_residual == idx_cpu.by_residual)

    idx_cpu.train(to_train)
    idx_cpu.add(to_train)

    res = faiss.StandardGpuResources()
    res.noTempMemory()
    idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist,
                                               qtype, metric, by_residual)
    assert(by_residual == idx_gpu.by_residual)

    idx_gpu.train(to_train)
    idx_gpu.add(to_train)

    return idx_cpu, idx_gpu

#
# Testing functions
#

def summarize_results(dist, idx):
    valid = []
    invalid = []
    for query in range(dist.shape[0]):
        valid_sub = {}
        invalid_sub = []

        for order, (d, i) in enumerate(zip(dist[query], idx[query])):
            if i == -1:
                invalid_sub.append(order)
            else:
                valid_sub[i] = [order, d]

        valid.append(valid_sub)
        invalid.append(invalid_sub)

    return valid, invalid

def compare_results(d1, i1, d2, i2):
    # Count number of index differences
    idx_diffs = {}
    idx_diffs_inf = 0
    idx_invalid = 0

    valid1, invalid1 = summarize_results(d1, i1)
    valid2, invalid2 = summarize_results(d2, i2)

    # Invalid results should be the same for both
    # (except if we happen to hit different centroids)
    for inv1, inv2 in zip(invalid1, invalid2):
        if (len(inv1) != len(inv2)):
            print('mismatch ', len(inv1), len(inv2), inv2[0])

        assert(len(inv1) == len(inv2))
        idx_invalid += len(inv2)
        for x1, x2 in zip(inv1, inv2):
            assert(x1 == x2)

    for _, (query1, query2) in enumerate(zip(valid1, valid2)):
        for idx1, order_d1 in query1.items():
            order_d2 = query2.get(idx1, None)
            if order_d2:
                idx_diff = order_d1[0] - order_d2[0]

                if idx_diff not in idx_diffs:
                    idx_diffs[idx_diff] = 1
                else:
                    idx_diffs[idx_diff] += 1
            else:
                idx_diffs_inf += 1

    return idx_diffs, idx_diffs_inf, idx_invalid

def check_diffs(total_num, in_window_thresh, diffs, diff_inf, invalid):
    # We require a certain fraction of results to be within +/- diff_window
    # index differences
    diff_window = 4
    in_window = 0

    for diff in sorted(diffs):
        if abs(diff) <= diff_window:
            in_window += diffs[diff] / total_num

    if (in_window < in_window_thresh):
        print('error {} {}'.format(in_window, in_window_thresh))
        assert(in_window >= in_window_thresh)

def do_test_with_index(ci, gi, nprobe, k, clamp, in_window_thresh):
    num_query = 11
    to_query = make_t(num_query, ci.d, clamp)

    ci.nprobe = ci.nprobe
    gi.nprobe = gi.nprobe

    total_num = num_query * k
    check_diffs(total_num, in_window_thresh,
                *compare_results(*ci.search(to_query, k),
                                 *gi.search(to_query, k)))

def do_test(nlist, d, qtype, by_residual, metric, nprobe, k):
    clamp = (qtype == faiss.ScalarQuantizer.QT_8bit_direct)
    ci, gi = make_indices_copy_from_cpu(nlist, d, qtype,
                                        by_residual, metric, clamp)
    # A direct copy should be much more closely in agreement
    # (except for fp accumulation order differences)
    do_test_with_index(ci, gi, nprobe, k, clamp, 0.99)

    ci, gi = make_indices_copy_from_gpu(nlist, d, qtype,
                                        by_residual, metric, clamp)
    # A direct copy should be much more closely in agreement
    # (except for fp accumulation order differences)
    do_test_with_index(ci, gi, nprobe, k, clamp, 0.99)

    ci, gi = make_indices_train(nlist, d, qtype,
                                by_residual, metric, clamp)
    # Separate training can produce a slightly different coarse quantizer
    # and residuals
    do_test_with_index(ci, gi, nprobe, k, clamp, 0.8)

def do_multi_test(qtype):
    nlist = 100
    nprobe = 10
    k = 50

    for d in [11, 64]:
        if (qtype != faiss.ScalarQuantizer.QT_8bit_direct):
            # residual doesn't make sense here
            do_test(nlist, d, qtype, True,
                    faiss.METRIC_L2, nprobe, k)
            do_test(nlist, d, qtype, True,
                    faiss.METRIC_INNER_PRODUCT, nprobe, k)
        do_test(nlist, d, qtype, False, faiss.METRIC_L2, nprobe, k)
        do_test(nlist, d, qtype, False, faiss.METRIC_INNER_PRODUCT, nprobe, k)

#
# Test
#

class TestSQ(unittest.TestCase):
    def test_fp16(self):
        do_multi_test(faiss.ScalarQuantizer.QT_fp16)

    def test_8bit(self):
        do_multi_test(faiss.ScalarQuantizer.QT_8bit)

    def test_8bit_uniform(self):
        do_multi_test(faiss.ScalarQuantizer.QT_8bit_uniform)

    def test_6bit(self):
        try:
            do_multi_test(faiss.ScalarQuantizer.QT_6bit)
            # should not reach here; QT_6bit is unimplemented
        except:
            print('QT_6bit exception thrown (is expected)')
        else:
            assert(False)

    def test_4bit(self):
        do_multi_test(faiss.ScalarQuantizer.QT_4bit)

    def test_4bit_uniform(self):
        do_multi_test(faiss.ScalarQuantizer.QT_4bit_uniform)

    def test_8bit_direct(self):
        do_multi_test(faiss.ScalarQuantizer.QT_8bit_direct)


if __name__ == '__main__':
    unittest.main()