image_output_1.png
from tensorflow.keras import Model

from tensorflow.keras.layers import Input, Dense

import matplotlib.pyplot as plt

import tensorflow as tf

def sample_dataset():
    dataset_shape = (2000, 1)
    return tf.random.normal(mean=8.0, shape=dataset_shape, stddev=0.5, dtype=tf.float32)
axes = plt.gca()
axes.set_xlim([-1, 11])
axes.set_ylim([0, 70])

def generator(input_shape):
    """Defines the generator keras.Model.
    Args:
        input_shape: the desired input shape (e.g.: (latent_space_size))
    Returns:
        G: The generator model
    """
    inputs = Input(input_shape)
    net = Dense(units=64, activation=tf.nn.elu, name='fc1')(inputs)
    net = Dense(units=64, activation=tf.nn.elu, name='fc2')(net)
    net = Dense(units=1, name='G')(net)
    G = Model(inputs=inputs, outputs=net)
    return G
def disciminator(input_shape):
    """Defines the Discriminator keras.Model.
    Args:
        input_shape: the desired input shape (e.g.: (the generator output shape))
    Returns:
        D: the Discriminator model
    """
    inputs = Input(input_shape)
    net = Dense(units=32, activation=tf.nn.elu, name='fc1')(inputs)
    net = Dense(units=1, name='D')(net)
    D = Model(inputs=inputs, outputs=net)
    return D

input_shape = (1,)
D = disciminator(input_shape)
latent_space_shape = (100,)
G = generator(latent_space_shape)

bce = tf.keras.losses.BinaryCrossentropy(from_logits=True)
def d_loss(d_real, d_fake):
    """The disciminator loss function."""
    return bce(tf.ones_like(d_real), d_real) + bce(tf.zeros_like(d_fake), d_fake)
def g_loss(generated_output):
    """The Generator loss function."""
    return bce(tf.ones_like(generated_output), generated_output)

optimizer = tf.keras.optimizers.Adam(1e-05)
@tf.function
def train_step():
    with tf.GradientTape(persistent=True) as tape:
        real_data = sample_dataset()
        noise_vector = tf.random.normal(mean=0, stddev=1, shape=(real_data.shape[0], latent_space_shape[0]))
        fake_data = G(noise_vector)
        d_fake_data = D(fake_data)
        d_real_data = D(real_data)
        d_loss_value = d_loss(d_real_data, d_fake_data)
        g_loss_value = g_loss(d_fake_data)
    d_gradients = tape.gradient(d_loss_value, D.trainable_variables)
    g_gradients = tape.gradient(g_loss_value, G.trainable_variables)
    del tape
    optimizer.apply_gradients(zip(d_gradients, D.trainable_variables))
    optimizer.apply_gradients(zip(g_gradients, G.trainable_variables))
    return (real_data, fake_data, g_loss_value, d_loss_value)

fig, ax = plt.subplots()
for step in range(40000):
    real_data, fake_data, g_loss_value, d_loss_value = train_step()
    if step % 2000 == 0:
        print('G loss: ', g_loss_value.numpy(), ' D loss: ', d_loss_value.numpy(), ' step: ', step)
        ax.hist(fake_data.numpy(), 100)
        ax.hist(real_data.numpy(), 100)
        props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
        textstr = f'step={step}'
        ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=14, verticalalignment='top', bbox=props)
        axes = plt.gca()
        axes.set_xlim([-1, 11])
        axes.set_ylim([0, 60])
        display(plt.gcf())
        plt.gca().clear()