Function learn_bopq

experiments/python/product_quantize.py:343–403

(X_train, ncodebooks, codebook_bits=4, niters=20,
               initial_kmeans_iters=1, R_sz=16, **sink)

@_memory.cache  # opq with block diagonal rotations
def learn_bopq(X_train, ncodebooks, codebook_bits=4, niters=20,
               initial_kmeans_iters=1, R_sz=16, **sink):

    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

    # compute number of rotations and subspaces associated with each
    nrots = int(D / R_sz)
    rot_starts = R_sz * np.arange(nrots)
    rot_ends = rot_starts + R_sz

    # X_rotated, R = opq_initialize(X_train, ncodebooks=ncodebooks, init=init)
    X_rotated = X  # hardcode identity init # TODO allow others
    rotations = [np.eye(R_sz) for i in range(nrots)]

    # 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 "---- BOPQ {} {}x{}b iter {}: mse / variance = {:.5f}".format(
            R_sz, ncodebooks, codebook_bits, it, err)

        rotations = []
        for i in range(nrots):
            start, end = rot_starts[i], rot_ends[i]

            X_sub = X[:, start:end]
            X_hat_sub = X_hat[:, start:end]

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

            X_rotated[:, start:end] = np.dot(X_sub, R.T)

        # update assignments and codebooks based on new rotations
        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 "---- BOPQ {} {}x{}b final mse / variance = {:.5f} ({:.3f}s)".format(