Convolutional variational autoencoder with PyMC3 and Keras¶

from keras.models import Sequential
from keras.layers import Dense, BatchNormalization

def get_params(model):
    """Get parameters and updates from Keras model
    """
    shared_in_updates = list()
    params = list()
    updates = dict()

    for l in model.layers:
        attrs = dir(l)
        # Updates
        if 'updates' in attrs:
            updates.update(l.updates)
            shared_in_updates += [e[0] for e in l.updates]

        # Shared variables
        for attr_str in attrs:
            attr = getattr(l, attr_str)
            if type(attr) is tt.sharedvar.TensorSharedVariable:
                if attr is not model.get_input_at(0):
                    params.append(attr)

    return list(set(params) - set(shared_in_updates)), updates

# This code is required when using BatchNormalization layer
keras.backend.theano_backend._LEARNING_PHASE = \
    shared(np.uint8(1), name='keras_learning_phase')

def cnn_enc(xs, latent_dim, nb_filters=64, nb_conv=3, intermediate_dim=128):
    """Returns a CNN model of Keras.

    Parameters
    ----------
    xs : theano.tensor.sharedvar.TensorSharedVariable
        Input tensor.
    latent_dim : int
        Dimension of latent vector.
    """
    input_layer = InputLayer(input_tensor=xs,
                             batch_input_shape=xs.get_value().shape)
    model = Sequential()
    model.add(input_layer)

    cp1 = {'border_mode': 'same', 'activation': 'relu'}
    cp2 = {'border_mode': 'same', 'activation': 'relu', 'subsample': (2, 2)}
    cp3 = {'border_mode': 'same', 'activation': 'relu', 'subsample': (1, 1)}
    cp4 = cp3

    model.add(Convolution2D(1, 2, 2, **cp1))
    model.add(Convolution2D(nb_filters, 2, 2, **cp2))
    model.add(Convolution2D(nb_filters, nb_conv, nb_conv, **cp3))
    model.add(Convolution2D(nb_filters, nb_conv, nb_conv, **cp4))
    model.add(Flatten())
    model.add(Dense(intermediate_dim, activation='relu'))
    model.add(Dense(2 * latent_dim))

    return model

class Encoder:
    """Encode observed images to variational parameters (mean/std of Gaussian).

    Parameters
    ----------
    xs : theano.tensor.sharedvar.TensorSharedVariable
        Placeholder of input images.
    dim_hidden : int
        The number of hidden variables.
    net : Function
        Returns
    """
    def __init__(self, xs, dim_hidden, net):
        model = net(xs, dim_hidden)

        self.model = model
        self.xs = xs
        self.out = model.get_output_at(-1)
        self.means = self.out[:, :dim_hidden]
        self.lstds = self.out[:, dim_hidden:]
        self.params, self.updates = get_params(model)
        self.enc_func = None
        self.dim_hidden = dim_hidden

    def _get_enc_func(self):
        if self.enc_func is None:
            xs = tt.tensor4()
            means = clone(self.means, {self.xs: xs})
            lstds = clone(self.lstds, {self.xs: xs})
            self.enc_func = function([xs], [means, lstds])

        return self.enc_func

    def encode(self, xs):
        # Used in test phase
        keras.backend.theano_backend._LEARNING_PHASE.set_value(np.uint8(0))

        enc_func = self._get_enc_func()
        means, _ = enc_func(xs)

        return means

    def draw_samples(self, xs, n_samples=1):
        """Draw samples of hidden variables based on variational parameters encoded.

        Parameters
        ----------
        xs : numpy.ndarray, shape=(n_images, 1, height, width)
            Images.
        """
        # Used in test phase
        keras.backend.theano_backend._LEARNING_PHASE.set_value(np.uint8(0))

        enc_func = self._get_enc_func()
        means, lstds = enc_func(xs)
        means = np.repeat(means, n_samples, axis=0)
        lstds = np.repeat(lstds, n_samples, axis=0)
        ns = np.random.randn(len(xs) * n_samples, self.dim_hidden)
        zs = means + np.exp(lstds) * ns

        return ns

def cnn_dec(zs, minibatch_size, nb_filters=64, nb_conv=3, dim_hidden=100, output_shape=(1, 28, 28)):
    """Returns a CNN model of Keras.

    Parameters
    ----------
    zs : theano.tensor.var.TensorVariable
        Input tensor.
    """
    input_layer = InputLayer(input_tensor=zs,
                             batch_input_shape=zs.tag.test_value.shape)
    model = Sequential()
    model.add(input_layer)

    model.add(Dense(dim_hidden, activation='relu'))
    model.add(Dense(nb_filters * 14 * 14, activation='relu'))

    cp1 = {'border_mode': 'same', 'activation': 'relu', 'subsample': (1, 1)}
    cp2 = cp1
    cp3 = {'border_mode': 'valid', 'activation': 'relu', 'subsample': (2, 2)}
    cp4 = {'border_mode': 'valid', 'activation': 'sigmoid'}

    output_shape_ = (minibatch_size, nb_filters, 14, 14)
    model.add(Reshape(output_shape_[1:]))
    model.add(Deconvolution2D(nb_filters, nb_conv, nb_conv, output_shape_, **cp1))
    model.add(Deconvolution2D(nb_filters, nb_conv, nb_conv, output_shape_, **cp2))
    output_shape_ = (minibatch_size, nb_filters, 29, 29)
    model.add(Deconvolution2D(nb_filters, 2, 2, output_shape_, **cp3))
    model.add(Convolution2D(1, 2, 2, **cp4))

    return model

class Decoder:
    """Decode hidden variables to images.

    Parameters
    ----------
    zs : Theano tensor
        Hidden variables.
    """
    def __init__(self, zs, net):
        model = net(zs, dim_hidden)
        self.model = model
        self.zs = zs
        self.out = model.get_output_at(-1)
        self.params, self.updates = get_params(model)
        self.dec_func = None

    def _get_dec_func(self):
        if self.dec_func is None:
            zs = tt.matrix()
            xs = clone(self.out, {self.zs: zs})
            self.dec_func = function([zs], xs)

        return self.dec_func

    def decode(self, zs):
        """Decode hidden variables to images.

        An image consists of the mean parameters of the observation noise.

        Parameters
        ----------
        zs : numpy.ndarray, shape=(n_samples, dim_hidden)
            Hidden variables.
        """
        # Used in test phase
        keras.backend.theano_backend._LEARNING_PHASE.set_value(np.uint8(0))

        return self._get_dec_func()(zs)

# Constants
minibatch_size = 100
dim_hidden = 2

with pm.Model() as model:
    # Hidden variables
    zs = pm.Normal('zs', mu=0, sd=1, shape=(minibatch_size, dim_hidden), dtype='float32')

    # Decoder and its parameters
    dec = Decoder(zs, net=cnn_dec)

    # Observation model
    xs_ = pm.Normal('xs_', mu=dec.out.ravel(), sd=1, observed=xs_t.ravel(), dtype='float32')

local_RVs = OrderedDict({zs: ((enc.means, enc.lstds), len(data) / float(minibatch_size))})

observed_RVs = OrderedDict({xs_: len(data) / float(minibatch_size)})

# Mini-batches
def create_minibatch(data, minibatch_size):
    rng = np.random.RandomState(0)
    start_idx = 0
    while True:
        # Return random data samples of set size batchsize each iteration
        ixs = rng.randint(data.shape[0], size=minibatch_size)
        yield data[ixs]

minibatches = zip(create_minibatch(data, minibatch_size))

def adam(loss, param):
    adam_ = keras.optimizers.Adam()
    return adam_.get_updates(param, [], loss)