import numpy as np from scipy.stats import norm import tensorflow as tf from matplotlib import pyplot as plt class Generator(tf.keras.Model): def __init__(self): super(Generator, self).__init__() self.w = tf.Variable(1, dtype=tf.float32) self.b = tf.Variable(0, dtype=tf.float32) def call(self, z): x = self.w*z + self.b return x class Discriminator(tf.keras.Model): def __init__(self, hidden_units=8): super(Discriminator, self).__init__() self.dense = tf.keras.layers.Dense(hidden_units) self.logits = tf.keras.layers.Dense(1, kernel_initializer=tf.keras.initializers.zeros()) def call(self, x): x = tf.expand_dims(x, axis=-1) logits = self.logits(tf.nn.relu(self.dense(x))) logits = tf.squeeze(logits, axis=-1) p = 1 / (1 + tf.math.exp(-logits)) return p # parameters true distribution mu = -4 sigma = 2 def visualize(G, D): interval = np.linspace(-10, 10, 100).astype(np.float32) d_values = D(interval) g_dist = norm.pdf(interval, loc=G.b.numpy(), scale=G.w.numpy()) true_dist = norm.pdf(interval, loc=mu, scale=sigma) plt.plot(interval, true_dist, label="true_dist") plt.plot(interval, g_dist, label="G_dist") plt.plot(interval, d_values, label="D, p_true_data(x)") plt.legend() plt.show() N = 32 x = np.random.normal(loc=mu, scale=sigma, size=N).astype(np.float32) indices = np.array(range(N)) G = Generator() D = Discriminator() learning_rate = 0.1 D_optimizer = tf.keras.optimizers.Adam(learning_rate) G_optimizer = tf.keras.optimizers.Adam(learning_rate) batch_size = 16 critic_iters = 2 iterations = 100 plot_interval = 1 for iteration in range(iterations): if iteration % plot_interval == 0: visualize(G, D) # discriminator update for _ in range(critic_iters): # sample real data (from data distribution) np.random.shuffle(indices) real = x[indices[:batch_size]] z = tf.random.normal(shape=[batch_size]) fake = G(z) with tf.GradientTape() as tape: loss_real = tf.reduce_mean(-tf.math.log(D(real))) loss_fake = tf.reduce_mean(-tf.math.log(1-D(fake))) D_loss = loss_real + loss_fake grads = tape.gradient(D_loss, D.trainable_variables) D_optimizer.apply_gradients(zip(grads, D.trainable_variables)) # generator update z = tf.random.normal(shape=[batch_size]) with tf.GradientTape() as tape: G_loss = tf.reduce_mean(tf.math.log(1-D(G(z)))) grads = tape.gradient(G_loss, G.trainable_variables) G_optimizer.apply_gradients(zip(grads, G.trainable_variables))