Function learn_opq

experiments/python/product_quantize.py:272–319

init in {'gauss', 'identity', 'random'}

(X_train, ncodebooks, codebook_bits=8, niters=20,
              initial_kmeans_iters=1, init='gauss', debug=False)

# https://github.com/arbabenko/Quantizations/blob/master/opqCoding.py
@_memory.cache
def learn_opq(X_train, ncodebooks, codebook_bits=8, niters=20,
              initial_kmeans_iters=1, init='gauss', debug=False):
    """init in {'gauss', 'identity', 'random'}"""

    print "OPQ: Using init '{}'".format(init)

    t0 = time.time()

    X = X_train.astype(np.float32)
    N, D = X.shape
    ncentroids = int(2**codebook_bits)
    subvect_len = D // ncodebooks

    assert D % subvect_len == 0  # equal number of dims for each codebook

    X_rotated, R = opq_initialize(X_train, ncodebooks=ncodebooks, init=init)

    # initialize codebooks by running kmeans on each rotated dim; this way,
    # setting niters=0 corresponds to normal PQ
    codebooks, assignments = learn_pq(X_rotated, ncentroids=ncentroids,
                                      nsubvects=ncodebooks,
                                      subvect_len=subvect_len,
                                      max_kmeans_iters=1)

    for it in np.arange(niters):
        # compute reconstruction errors
        X_hat = reconstruct_X_pq(assignments, codebooks)
        # err = compute_reconstruction_error(X_rotated, X_hat, subvect_len=subvect_len)
        err = compute_reconstruction_error(X_rotated, X_hat)
        print "---- OPQ {}x{}b iter {}: mse / variance = {:.5f}".format(
            ncodebooks, codebook_bits, it, err)

        # update rotation matrix based on reconstruction errors
        U, s, V = np.linalg.svd(np.dot(X_hat.T, X), full_matrices=False)
        R = np.dot(U, V)

        # update centroids using new rotation matrix
        X_rotated = opq_rotate(X, R)
        assignments = _encode_X_pq(X_rotated, codebooks)
        codebooks = _update_centroids_opq(X_rotated, assignments, ncentroids)

    X_hat = reconstruct_X_pq(assignments, codebooks)
    err = compute_reconstruction_error(X_rotated, X_hat)
    t = time.time() - t0
    print "---- OPQ {}x{}b final mse / variance = {:.5f} ({:.3f}s)".format(
        ncodebooks, codebook_bits, err, t)

    return codebooks, assignments, R


# ================================================================ Block OPQ