Function opq_initialize

experiments/python/product_quantize.py:231–266

(X_train, ncodebooks, init='gauss')

# @_memory.cache
def opq_initialize(X_train, ncodebooks, init='gauss'):
    X = X_train
    _, D = X.shape

    if init == 'gauss' or init == 'gauss_flat' or init == 'gauss_shuffle':
        permute = (init == 'gauss_shuffle')
        R = learn_opq_gaussian_rotation(X_train, ncodebooks, shuffle=permute)
        R = R.astype(np.float32)

        if init == 'gauss_flat':
            # assert R.shape[0] == R.shape[1]
            D = R.shape[1]
            d = D / ncodebooks
            assert int(d) == d  # need same number of dims in each subspace
            # d = int(d)
            local_r = random_rotation(int(d))
            tiled = np.zeros((D, D))
            for c in range(ncodebooks):
                start = c * d
                end = start + d
                tiled[start:end, start:end] = local_r

            R = np.dot(R, tiled)

        X_rotated = opq_rotate(X, R)
    elif init == 'identity':
        R = np.identity(D, dtype=np.float32)  # D x D
        X_rotated = X
    elif init == 'random':
        R = np.random.randn(D, D).astype(np.float32)
        R = orthonormalize_rows(R)
        X_rotated = opq_rotate(X, R)
    else:
        raise ValueError("Unrecognized initialization method: ".format(init))

    return X_rotated, R


# loosely based on: