debug_th.py

import theano as th
import theano.tensor as T
import lasagne
import numpy as np
from theano.tensor.shared_randomstreams import RandomStreams
from lasagne.nonlinearities import tanh
import matplotlib.pyplot as plt
import sys

# change to True and see what happens
use_reshape = False
th.config.optimizer = 'None'

D = 28 * 28

mnist = np.array(np.load('mnist.npz')['X'], dtype=np.float32)
np.divide(mnist, mnist.max(), mnist)

input_data = T.matrix('X')

batch = 1000
H = 5

rec = lasagne.layers.InputLayer((batch, D), input_data)
rec = lasagne.layers.DenseLayer(rec, 500,
                                W=lasagne.init.Normal(), nonlinearity=tanh)
rec = lasagne.layers.DenseLayer(rec, 2 * H,
                                W=lasagne.init.Normal(), nonlinearity=None)
phi = lasagne.layers.get_output(rec)

if use_reshape:
    phi = T.reshape(phi, (2, batch, H))
    zmu = T.reshape(phi[0, :, :], (batch, H))
    zsigma = T.reshape(T.exp(phi[1, :, :]), (batch, H))
else:
    zmu = phi[:, :H]
    zsigma = T.exp(phi[:, H:])

rng = RandomStreams()
eps = rng.normal((batch, H))

z = eps * zsigma + zmu

gen = lasagne.layers.InputLayer((batch, H), z)
gen = lasagne.layers.DenseLayer(gen, 500,
                                W=lasagne.init.Normal(), nonlinearity=tanh)
gen = lasagne.layers.DenseLayer(gen, D,
                                W=lasagne.init.Normal(), nonlinearity=lasagne.nonlinearities.sigmoid)

obs = T.reshape(lasagne.layers.get_output(gen), (batch, D))


def bernoulli_density(x, p):
    return T.sum(x * T.log(p) + (1. - x) * T.log(1. - p), axis=1)

pmu = T.zeros((batch, H))
psigma = T.ones((batch, H))


def kl(mu1, sigma1, mu2, sigma2):
    ratio = T.square(sigma1 / sigma2)
    return 0.5 * T.sum((T.square((mu1 - mu2) / sigma2) + (ratio - 1 - T.log(ratio))), axis=1)

lower_bound = T.mean(bernoulli_density(input_data, obs) - kl(zmu, zsigma, pmu, psigma))

params = lasagne.layers.get_all_params(gen, trainable=True) + lasagne.layers.get_all_params(rec, trainable=True)
lr = T.scalar()

objective = -lower_bound

updates = lasagne.updates.adam(objective, params, learning_rate=lr)

train_fn = th.function([input_data, lr], [lower_bound, obs[:5, :]], updates=updates,
                       allow_input_downcast=True)

avg_lb = 0.
lr = 1e-4
for i in xrange(10000):
    perm = np.random.choice(mnist.shape[0], batch)
    data = mnist[perm]

    lb, sample = train_fn(mnist[perm], lr)
    avg_lb += 0.01 * (lb - avg_lb)

    sys.stdout.write("\r%f %f %d %f" % (avg_lb, lb, i, lr))

    if np.random.rand() < 0.3:
        x_samples = np.hstack((sample.reshape((28 * 5, 28)), data[:5].reshape(28 * 5, 28)))
        plt.matshow(x_samples)
        plt.savefig('samples.png')
        plt.close()