ml_code/wta_vae/utils.py at master · mschmidt87/ml_code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data


def load_data(train_examples, validation_examples):
    """
    Load MNIST data and subsample to the specified number of
    training and validation examples.

    Parameters
    ----------
    train_examples : int
        Number of training examples.
    validation_examples : int
        Number of validation examples.

    Returns
    -------
    mnist : tensorflow dataset
        The complete MNIST dataset
    train_images, train_labels : tuple of numpy.ndarrays
        Images and labels for training
    validation_images, validation_labels : tuple of numpy.ndarrays
        Images and labels for validation
    """
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
    train_images, train_labels = (mnist.train.images[:train_examples],
                                  mnist.train.labels[:train_examples])
    validation_images, validation_labels = (mnist.validation.images[:validation_examples],
                                            mnist.validation.labels[:validation_examples])
    return mnist, (train_images, train_labels), (validation_images, validation_labels)


def loss_function(x_recon, x, mean, logvar, batch_size):
    """
    Loss function of the Variational Autoencoder.
    Uses the mean-squared error for the reconstruction error.

    Parameters
    ----------
    x_recon : tensorflow.Variable
        The reconstructed input from the autoencoder.
    x : tensorflow.Variable
        The input data to the autoencoder
    mean : tensorflow.Variable
        The mean values of the latent variables
    logvar : tensorflow.Variable
        The logarithmic variance of the latent variables
    batch_size : int
        The size of the batch

    Returns
    -------
    loss : tensorflow.Variable
        The value of the loss function
    """

    MSE = tf.losses.mean_squared_error(x_recon, x)
    KLD = -0.5 * tf.reduce_sum(1 + logvar - tf.square(mean) - tf.exp(logvar))
    KLD /= batch_size * 784
    loss = MSE + KLD
    return loss