gmm.py

import argparse
import contextlib
import time

import matplotlib.pyplot as plt
import numpy as np
from scipy import stats

import cupy


@contextlib.contextmanager
def timer(message):
    cupy.cuda.Stream.null.synchronize()
    start = time.time()
    yield
    cupy.cuda.Stream.null.synchronize()
    end = time.time()
    print('%s:  %f sec' % (message, end - start))


def estimate_log_prob(X, inv_cov, means):
    xp = cupy.get_array_module(X)
    n_features = X.shape[1]
    log_det = xp.sum(xp.log(inv_cov), axis=1)
    precisions = inv_cov ** 2
    log_prob = xp.sum((means ** 2 * precisions), 1) - \
        2 * xp.dot(X, (means * precisions).T) + xp.dot(X ** 2, precisions.T)
    return -0.5 * (n_features * xp.log(2 * np.pi) + log_prob) + log_det


def m_step(X, resp):
    xp = cupy.get_array_module(X)
    nk = xp.sum(resp, axis=0)
    means = xp.dot(resp.T, X) / nk[:, None]
    X2 = xp.dot(resp.T, X * X) / nk[:, None]
    covariances = X2 - means ** 2
    return nk / len(X), means, covariances


def e_step(X, inv_cov, means, weights):
    xp = cupy.get_array_module(X)
    weighted_log_prob = estimate_log_prob(X, inv_cov, means) + \
        xp.log(weights)
    log_prob_norm = xp.log(xp.sum(xp.exp(weighted_log_prob), axis=1))
    log_resp = weighted_log_prob - log_prob_norm[:, None]
    return xp.mean(log_prob_norm), log_resp


def train_gmm(X, max_iter, tol, means, covariances):
    xp = cupy.get_array_module(X)
    lower_bound = -np.inf
    converged = False
    weights = xp.array([0.5, 0.5], dtype=np.float32)
    inv_cov = 1 / xp.sqrt(covariances)

    for n_iter in range(max_iter):
        prev_lower_bound = lower_bound
        log_prob_norm, log_resp = e_step(X, inv_cov, means, weights)
        weights, means, covariances = m_step(X, xp.exp(log_resp))
        inv_cov = 1 / xp.sqrt(covariances)
        lower_bound = log_prob_norm
        change = lower_bound - prev_lower_bound
        if abs(change) < tol:
            converged = True
            break

    if not converged:
        print('Failed to converge. Increase max-iter or tol.')

    return inv_cov, means, weights, covariances


def predict(X, inv_cov, means, weights):
    xp = cupy.get_array_module(X)
    log_prob = estimate_log_prob(X, inv_cov, means)
    return (log_prob + xp.log(weights)).argmax(axis=1)


def calc_acc(X_train, y_train, X_test, y_test, max_iter, tol, means,
             covariances):
    xp = cupy.get_array_module(X_train)
    inv_cov, means, weights, cov = \
        train_gmm(X_train, max_iter, tol, means, covariances)
    y_train_pred = predict(X_train, inv_cov, means, weights)
    train_accuracy = xp.mean(y_train_pred == y_train) * 100
    y_test_pred = predict(X_test, inv_cov, means, weights)
    test_accuracy = xp.mean(y_test_pred == y_test) * 100
    print('train_accuracy : %f' % train_accuracy)
    print('test_accuracy : %f' % test_accuracy)
    return y_test_pred, means, cov


def draw(X, pred, means, covariances, output):
    xp = cupy.get_array_module(X)
    for i in range(2):
        labels = X[pred == i]
        if xp is cupy:
            labels = labels.get()
        plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(1, 3))
    if xp is cupy:
        means = means.get()
        covariances = covariances.get()
    plt.scatter(means[:, 0], means[:, 1], s=120, marker='s', facecolors='y',
                edgecolors='k')
    x = np.linspace(-5, 5, 1000)
    y = np.linspace(-5, 5, 1000)
    X, Y = np.meshgrid(x, y)
    for i in range(2):
        dist = stats.multivariate_normal(means[i], covariances[i])
        Z = dist.pdf(np.stack([X, Y], axis=-1))
        plt.contour(X, Y, Z)
    plt.savefig(output)


def run(gpuid, num, dim, max_iter, tol, output):
    """CuPy Gaussian Mixture Model example

    Compute GMM parameters, weights, means and covariance matrix, depending on
    sampled data. There are two main components, e_step and m_step.
    In e_step, compute burden rate, which is expressed `resp`, by latest
    weights, means and covariance matrix.
    In m_step, compute weights, means and covariance matrix by latest `resp`.

    """
    scale = np.ones(dim)
    train1 = np.random.normal(1, scale, size=(num, dim)).astype(np.float32)
    train2 = np.random.normal(-1, scale, size=(num, dim)).astype(np.float32)
    X_train = np.r_[train1, train2]
    test1 = np.random.normal(1, scale, size=(100, dim)).astype(np.float32)
    test2 = np.random.normal(-1, scale, size=(100, dim)).astype(np.float32)
    X_test = np.r_[test1, test2]
    y_train = np.r_[np.zeros(num), np.ones(num)].astype(np.int32)
    y_test = np.r_[np.zeros(100), np.ones(100)].astype(np.int32)

    mean1 = np.random.normal(1, scale, size=dim)
    mean2 = np.random.normal(-1, scale, size=dim)
    means = np.stack([mean1, mean2])
    covariances = np.random.rand(2, dim)
    print('Running CPU...')
    with timer(' CPU '):
        y_test_pred, means, cov = \
            calc_acc(X_train, y_train, X_test, y_test, max_iter, tol,
                     means, covariances)

    with cupy.cuda.Device(gpuid):
        X_train_gpu = cupy.array(X_train)
        y_train_gpu = cupy.array(y_train)
        y_test_gpu = cupy.array(y_test)
        X_test_gpu = cupy.array(X_test)
        means = cupy.array(means)
        covariances = cupy.array(covariances)
        print('Running GPU...')
        with timer(' GPU '):
            y_test_pred, means, cov = \
                calc_acc(X_train_gpu, y_train_gpu, X_test_gpu, y_test_gpu,
                         max_iter, tol, means, covariances)
        if output is not None:
            draw(X_test_gpu, y_test_pred, means, cov, output)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu-id', '-g', default=0, type=int,
                        help='ID of GPU.')
    parser.add_argument('--num', '-n', default=500000, type=int,
                        help='number of train data')
    parser.add_argument('--dim', '-d', default=2, type=int,
                        help='dimension of each data')
    parser.add_argument('--max-iter', '-m', default=30, type=int,
                        help='number of iterations')
    parser.add_argument('--tol', '-t', default=1e-3, type=float,
                        help='error tolerance to stop iterations')
    parser.add_argument('--output-image', '-o', default=None, type=str,
                        dest='output', help='output image file name')
    args = parser.parse_args()
    run(args.gpu_id, args.num, args.dim, args.max_iter, args.tol, args.output)