autoencoder_test.py

# Autoencoder example for dimensionalilty reduction of MNIST
# Requires tensorflow and seaborn
# e.g.: pip install tensorflow-cpu seaborn

import tensorflow as tf
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt


def do_autoencoder_things():
    print("TensorFlow version:", tf.__version__)

    # load the data set (mnist)
    mnist = tf.keras.datasets.mnist
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    # normalize pixel values to range [0,1]
    x_train, x_test = x_train / 255.0, x_test / 255.0

    # reshape the images to flat lines of pixels (from shape [n][28][28] to [n][28*28])
    x_train = x_train.reshape(x_train.shape[0], x_train.shape[1] * x_train.shape[2])
    x_test = x_test.reshape(x_test.shape[0], x_test.shape[1] * x_test.shape[2])

    # set up the layers of the network
    # going from 28*28=784 to 128 to 32 to 8 to 2 dimensions
    l1 = tf.keras.layers.Dense(128, activation='sigmoid')
    l2 = tf.keras.layers.Dense(32, activation='sigmoid')
    l3 = tf.keras.layers.Dense(8, activation='sigmoid')
    l4 = tf.keras.layers.Dense(2, activation='sigmoid')
    # now going from 2 dimensions up again
    l5 = tf.keras.layers.Dense(8, activation='relu')
    l6 = tf.keras.layers.Dense(32, activation='relu')
    l7 = tf.keras.layers.Dense(128, activation='relu')
    l8 = tf.keras.layers.Dense(x_train.shape[1], activation='relu')

    # creating an encoder model from our layers
    encoder = tf.keras.models.Sequential(
        [l1,l2,l3,l4])
    # creating a combined encoder and decoder model from our layers
    encoder_decoder = tf.keras.models.Sequential(
        [l1,l2,l3,l4,l5,l6,l7,l8])

    # we will train the encoder_decoder model, therefore we compile it
    encoder_decoder.compile(optimizer='adam',
                  loss=tf.keras.losses.MeanSquaredError(),
                  metrics=['mean_absolute_error'])
    # check model performance before training
    print(f'model evaluation: {encoder_decoder.evaluate(x_test, x_test, verbose=2)}')
    # train the model, input and ground truth are the same because we optimize for best reconstruction
    # (using x_test instead of x_train because it's smaller and training will be quicker)
    encoder_decoder.fit(x_test, x_test, epochs=50)
    # check perfromance after training
    print(f'model evaluation: {encoder_decoder.evaluate(x_test, x_test, verbose=2)}')

    # now that the weights in the layers are trained, we can use the encoder
    reduced = encoder.predict(x_test)

    # plot the resulting 2D points
    sns.scatterplot(x=reduced[:, 0], y=reduced[:, 1], hue=y_test, palette="Paired")
    plt.show()


def main():
    do_autoencoder_things()


if __name__ == '__main__':
    main()