unsupervised_aae_deterministic_w_discriminator.py

"""
Deterministic unsupervised adversarial autoencoder.

 We are using:
    - Gaussian distribution as prior distribution.
    - Convolutional layers.
    - Discriminator in x space
"""
import time
from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib import gridspec
import matplotlib.patches as mpatches
import numpy as np
import tensorflow as tf


PROJECT_ROOT = Path.cwd()

# -------------------------------------------------------------------------------------------------------------
# Set random seed
random_seed = 42
tf.random.set_seed(random_seed)
np.random.seed(random_seed)

# -------------------------------------------------------------------------------------------------------------
output_dir = PROJECT_ROOT / 'outputs'
output_dir.mkdir(exist_ok=True)

experiment_dir = output_dir / 'unsupervised_aae_deterministic_w_discriminator'
experiment_dir.mkdir(exist_ok=True)

latent_space_dir = experiment_dir / 'latent_space'
latent_space_dir.mkdir(exist_ok=True)

reconstruction_dir = experiment_dir / 'reconstruction'
reconstruction_dir.mkdir(exist_ok=True)

sampling_dir = experiment_dir / 'sampling'
sampling_dir.mkdir(exist_ok=True)

# -------------------------------------------------------------------------------------------------------------
# Loading data
print("Loading data...")
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.

x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)

# -------------------------------------------------------------------------------------------------------------
# Create the dataset iterator
batch_size = 256
train_buf = 60000

train_dataset = tf.data.Dataset.from_tensor_slices(x_train)
train_dataset = train_dataset.shuffle(buffer_size=train_buf)
train_dataset = train_dataset.batch(batch_size)


# -------------------------------------------------------------------------------------------------------------
# Create models
def make_encoder_model(z_size):
    inputs = tf.keras.layers.Input(shape=(28, 28, 1))

    x = tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=2, padding='same')(inputs)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Conv2D(filters=128, kernel_size=3, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    z = tf.keras.layers.Conv2D(filters=z_size, kernel_size=3, strides=2, padding='same')(x)

    model = tf.keras.Model(inputs=inputs, outputs=z)
    return model


def make_decoder_model(z_size):
    encoded = tf.keras.Input(shape=(1, 1, z_size))

    x = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(encoded)
    x = tf.keras.layers.UpSampling2D((2, 2))(x)
    x = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(x)
    x = tf.keras.layers.UpSampling2D((2, 2))(x)
    x = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(x)
    x = tf.keras.layers.UpSampling2D((2, 2))(x)
    x = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='relu')(x)
    x = tf.keras.layers.UpSampling2D((2, 2))(x)
    x = tf.keras.layers.Conv2D(64, kernel_size=3, activation='relu')(x)
    x = tf.keras.layers.UpSampling2D((2, 2))(x)

    reconstruction = tf.keras.layers.Conv2D(filters=1, kernel_size=3, activation='sigmoid', padding='same')(x)
    decoder = tf.keras.Model(inputs=encoded, outputs=reconstruction)
    return decoder



def make_discriminator_z_model(z_size):
    encoded = tf.keras.Input(shape=(z_size,))
    x = tf.keras.layers.Dense(128)(encoded)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Dense(128)(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    prediction = tf.keras.layers.Dense(1)(x)
    model = tf.keras.Model(inputs=encoded, outputs=prediction)
    return model


def make_discriminator_x_model():
    inputs = tf.keras.layers.Input(shape=(28, 28, 1))

    x = tf.keras.layers.Conv2D(filters=16, kernel_size=4, strides=2, padding='same')(inputs)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Conv2D(filters=32, kernel_size=4, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    x = tf.keras.layers.Conv2D(filters=64, kernel_size=4, strides=2, padding='same')(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.LeakyReLU(0.2)(x)
    z = tf.keras.layers.Conv2D(filters=1, kernel_size=4, strides=1, padding='valid')(x)

    model = tf.keras.Model(inputs=inputs, outputs=z)
    return model



z_dim = 2
encoder = make_encoder_model(z_dim)
decoder = make_decoder_model(z_dim)
discriminator_z = make_discriminator_z_model(z_dim)
discriminator_x = make_discriminator_x_model()


# -------------------------------------------------------------------------------------------------------------
# Define loss functions
ae_loss_weight = 1.
gen_loss_weight = 1.
dc_loss_weight = 1.

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
mse = tf.keras.losses.MeanSquaredError()
accuracy = tf.keras.metrics.BinaryAccuracy()


def autoencoder_loss(inputs, reconstruction, loss_weight):
    return loss_weight * mse(inputs, reconstruction)


def discriminator_loss(real_output, fake_output, loss_weight):
    loss_real = cross_entropy(tf.ones_like(real_output), real_output)
    loss_fake = cross_entropy(tf.zeros_like(fake_output), fake_output)
    return loss_weight * (loss_fake + loss_real)


def generator_loss(fake_output, loss_weight):
    return loss_weight * cross_entropy(tf.ones_like(fake_output), fake_output)


# -------------------------------------------------------------------------------------------------------------
# Define optimizers
learning_rate = 0.0001

ae_optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
dc_z_optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
gen_z_optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
dc_x_optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
gen_x_optimizer = tf.keras.optimizers.Adam(lr=learning_rate)


@tf.function
def train_step(batch_x):
    # -------------------------------------------------------------------------------------------------------------
    # Autoencoder
    with tf.GradientTape() as ae_tape:
        encoder_output = encoder(batch_x, training=True)
        decoder_output = decoder(encoder_output, training=True)

        # Autoencoder loss
        ae_loss = autoencoder_loss(batch_x, decoder_output, ae_loss_weight)

    ae_grads = ae_tape.gradient(ae_loss, encoder.trainable_variables + decoder.trainable_variables)
    ae_optimizer.apply_gradients(zip(ae_grads, encoder.trainable_variables + decoder.trainable_variables))


    # -------------------------------------------------------------------------------------------------------------
    # Discriminator Z
    with tf.GradientTape() as dc_tape:
        real_distribution = tf.random.normal([batch_size, 1, 1, z_dim], mean=0.0, stddev=1.0)
        encoder_output = encoder(batch_x, training=True)

        dc_z_real = discriminator_z(real_distribution, training=True)
        dc_z_fake = discriminator_z(encoder_output, training=True)

        # Discriminator Loss
        dc_z_loss = discriminator_loss(dc_z_real, dc_z_fake, dc_loss_weight)

        # Discriminator Acc
        dc_z_acc = accuracy(tf.concat([tf.ones_like(dc_z_real), tf.zeros_like(dc_z_fake)], axis=0),
                          tf.concat([dc_z_real, dc_z_fake], axis=0))

    dc_grads = dc_tape.gradient(dc_z_loss, discriminator_z.trainable_variables)
    # dc_z_optimizer.apply_gradients(zip(dc_grads, discriminator_z.trainable_variables))

    # -------------------------------------------------------------------------------------------------------------
    # Generator Z (Encoder)
    with tf.GradientTape() as gen_tape:
        encoder_output = encoder(batch_x, training=True)
        dc_z_fake = discriminator_z(encoder_output, training=True)

        # Generator loss
        gen_z_loss = generator_loss(dc_z_fake, gen_loss_weight)

    gen_z_grads = gen_tape.gradient(gen_z_loss, encoder.trainable_variables)
    # gen_z_optimizer.apply_gradients(zip(gen_z_grads, encoder.trainable_variables))

    # -------------------------------------------------------------------------------------------------------------
    # Discriminator X
    with tf.GradientTape() as dc_x_tape:
        encoder_output = encoder(batch_x, training=True)
        decoder_output = decoder(encoder_output, training=True)

        d_x_real = discriminator_x(batch_x, training=True)
        d_x_fake = discriminator_x(decoder_output, training=True)

        # Discriminator X Loss
        dc_x_loss = discriminator_loss(d_x_real, d_x_fake, dc_loss_weight)

        # Discriminator X Acc
        dc_z_acc = accuracy(tf.concat([tf.ones_like(d_x_real), tf.zeros_like(d_x_fake)], axis=0),
                          tf.concat([d_x_real, d_x_fake], axis=0))

    dc_x_grads = dc_x_tape.gradient(dc_x_loss, discriminator_x.trainable_variables)
    # dc_x_optimizer.apply_gradients(zip(dc_x_grads, discriminator_x.trainable_variables))

    # -------------------------------------------------------------------------------------------------------------
    # Generator X (Decoder)
    with tf.GradientTape() as gen_x_tape:
        encoder_output = encoder(batch_x, training=True)
        decoder_output = decoder(encoder_output, training=True)

        # Generator X loss
        d_x_fake = discriminator_x(decoder_output, training=True)

        gen_x_loss = generator_loss(d_x_fake, gen_loss_weight)

    gen_x_grads = gen_x_tape.gradient(gen_x_loss, decoder.trainable_variables)
    # gen_x_optimizer.apply_gradients(zip(gen_x_grads, decoder.trainable_variables))

    return ae_loss, dc_z_loss, dc_z_acc, gen_z_loss, dc_x_loss, dc_x_acc, gen_x_loss


# -------------------------------------------------------------------------------------------------------------
# Training loop
n_epochs = 200
for epoch in range(n_epochs):
    start = time.time()

    epoch_ae_loss_avg = tf.metrics.Mean()
    epoch_dc_z_loss_avg = tf.metrics.Mean()
    epoch_dc_z_acc_avg = tf.metrics.Mean()
    epoch_gen_z_loss_avg = tf.metrics.Mean()
    epoch_dc_x_loss_avg = tf.metrics.Mean()
    epoch_dc_x_acc_avg = tf.metrics.Mean()
    epoch_gen_x_loss_avg = tf.metrics.Mean()

    for batch, (batch_x) in enumerate(train_dataset):
        ae_loss, dc_z_loss, dc_z_acc, gen_z_loss, dc_x_loss, dc_x_acc, gen_x_loss = train_step(batch_x)

        epoch_ae_loss_avg(ae_loss)
        epoch_dc_z_loss_avg(dc_z_loss)
        epoch_dc_z_acc_avg(dc_z_acc)
        epoch_gen_z_loss_avg(gen_z_loss)
        epoch_dc_x_loss_avg(dc_x_loss)
        epoch_dc_x_acc_avg(dc_x_acc)
        epoch_gen_x_loss_avg(gen_x_loss)


    epoch_time = time.time() - start
    print(
        '{:4d}: TIME: {:.2f} ETA: {:.2f} AE_LOSS: {:.4f} DC_Z_LOSS: {:.4f} DC_Z_ACC: {:.4f} GEN_Z_LOSS: {:.4f} DC_X_LOSS: {:.4f} DC_X_ACC: {:.4f} GEN_X_LOSS: {:.4f}'
        .format(epoch, epoch_time,
                epoch_time * (n_epochs - epoch),
                epoch_ae_loss_avg.result(),
                epoch_dc_z_loss_avg.result(),
                epoch_dc_z_acc_avg.result(),
                epoch_gen_z_loss_avg.result(),
                epoch_dc_x_loss_avg.result(),
                epoch_dc_x_acc_avg.result(),
                epoch_gen_x_loss_avg.result()))


    # -------------------------------------------------------------------------------------------------------------
    if epoch % 10 == 0:
        # Latent Space
        x_test_encoded = encoder(x_test, training=False)
        label_list = list(y_test)

        fig = plt.figure()
        classes = set(label_list)
        colormap = plt.cm.rainbow(np.linspace(0, 1, len(classes)))
        kwargs = {'alpha': 0.8, 'c': [colormap[i] for i in label_list]}
        ax = plt.subplot(111, aspect='equal')
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
        handles = [mpatches.Circle((0, 0), label=class_, color=colormap[i])
                   for i, class_ in enumerate(classes)]
        ax.legend(handles=handles, shadow=True, bbox_to_anchor=(1.05, 0.45),
                  fancybox=True, loc='center left')
        plt.scatter(x_test_encoded[:, :, :, 0], x_test_encoded[:, :, :, 1], s=2, **kwargs)
        ax.set_xlim([-3, 3])
        ax.set_ylim([-3, 3])

        plt.savefig(latent_space_dir / ('epoch_%d.png' % epoch))
        plt.close('all')

        # Reconstruction
        n_digits = 20  # how many digits we will display
        x_test_decoded = decoder(encoder(x_test[:n_digits], training=False), training=False)
        x_test_decoded = np.reshape(x_test_decoded, [-1, 28, 28]) * 255
        fig = plt.figure(figsize=(20, 4))
        for i in range(n_digits):
            # display original
            ax = plt.subplot(2, n_digits, i + 1)
            plt.imshow(x_test[i].reshape(28, 28))
            plt.gray()
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)

            # display reconstruction
            ax = plt.subplot(2, n_digits, i + 1 + n_digits)
            plt.imshow(x_test_decoded[i])
            plt.gray()
            ax.get_xaxis().set_visible(False)
            ax.get_yaxis().set_visible(False)

        plt.savefig(reconstruction_dir / ('epoch_%d.png' % epoch))
        plt.close('all')

        # Sampling
        x_points = np.linspace(-3, 3, 20).astype(np.float32)
        y_points = np.linspace(-3, 3, 20).astype(np.float32)

        nx, ny = len(x_points), len(y_points)
        plt.subplot()
        gs = gridspec.GridSpec(nx, ny, hspace=0.05, wspace=0.05)

        for i, g in enumerate(gs):
            z = np.concatenate(([x_points[int(i / ny)]], [y_points[int(i % nx)]]))
            z = np.reshape(z, (1, 1, 1, 2))
            x = decoder(z, training=False).numpy()
            ax = plt.subplot(g)
            img = np.array(x.tolist()).reshape(28, 28)
            ax.imshow(img, cmap='gray')
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_aspect('auto')
        plt.savefig(sampling_dir / ('epoch_%d.png' % epoch))
        plt.close('all')