Class OPQEncoder

experiments/python/main.py:187–312

class OPQEncoder(PQEncoder):

    def __init__(self, dataset, code_bits=-1, bits_per_subvect=-1,
                 nsubvects=-1, elemwise_dist_func=dists_elemwise_sq,
                 opq_iters=20, quantize_lut=False, algo='OPQ', **opq_kwargs):
        X = dataset.X_train
        self.elemwise_dist_func = elemwise_dist_func
        self.quantize_lut = quantize_lut
        self.opq_iters = opq_iters
        self.algo = algo

        tmp = _parse_codebook_params(X.shape[1], code_bits=code_bits,
                                     bits_per_subvect=bits_per_subvect,
                                     nsubvects=nsubvects)
        self.nsubvects, self.ncentroids, self.subvect_len = tmp
        self.code_bits = int(np.log2(self.ncentroids))

        # for fast lookups via indexing into flattened array
        self.offsets = np.arange(self.nsubvects, dtype=np.int) * self.ncentroids

        if self.algo == 'Bolt':
            # Note: we always pass in 0 iters in the reported experiments,
            # so it never rotates anything
            self.centroids, _, self.rotations = pq.learn_bopq(
                X, ncodebooks=nsubvects, codebook_bits=bits_per_subvect,
                niters=opq_iters, **opq_kwargs)
        elif self.algo == 'OPQ':
            self.centroids, _, self.R = pq.learn_opq(
                X, ncodebooks=nsubvects, codebook_bits=bits_per_subvect,
                niters=opq_iters, **opq_kwargs)
        else:
            raise ValueError("argument algo must be one of {OPQ, Bolt}")

        # learn appropriate offsets and shared scale factor for quantization
        self.lut_offsets = np.zeros(self.nsubvects)
        self.order_idxs = np.arange(self.nsubvects, dtype=np.int)

        if self.quantize_lut:  # TODO put this logic in separate function
            print "learning quantization..."

            num_rows = min(10*1000, len(X) / 2)
            _, queries = datasets.extract_random_rows(
                X[num_rows:], how_many=1000, remove_from_X=False)
            X = X[:num_rows]  # limit to first 10k rows of X

            # compute luts for all the queries
            luts = [self._fit_query(q, quantize=False) for q in queries]
            luts = np.vstack(luts)
            assert luts.shape == (self.ncentroids * len(queries), self.nsubvects)

            self.lut_offsets, self.scale_by, _ = _learn_best_quantization(luts)

    def name(self):
        return "{}_{}x{}b_iters={}_quantize={}".format(
            self.algo, self.nsubvects, self.code_bits, self.opq_iters,
            int(self.quantize_lut))

    def params(self):