# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import tensorflow as tf
import numpy as np
import pandas as pd
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plot
from tensorflow.examples.tutorials.mnist import input_data
from optim.optimisticmirroradam import AdamirrorOptimizer

batch_size = 32
X_dim = 784
z_dim = 10
h_dim = 128
lambda_grad = 10
lr = 0.0001
n_disc = 5

mnist = input_data.read_data_sets('../../MNIST_data', one_hot=True)


def plot_images(samples):
    fig = plot.figure(figsize=(4, 4))
    gs = gridspec.GridSpec(4, 4)
    gs.update(wspace=0.05, hspace=0.05)

    for i, sample in enumerate(samples):
        ax = plot.subplot(gs[i])
        plot.axis('off')
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_aspect('equal')
        plot.imshow(sample.reshape(28, 28), cmap='Greys_r')

    return fig


# +
X = tf.placeholder(tf.float32, shape=[None, X_dim])

# D_W1 = tf.Variable(xavier_init([X_dim, h_dim]))
# D_b1 = tf.Variable(tf.zeros(shape=[h_dim]))

# D_W2 = tf.Variable(xavier_init([h_dim, 1]))
# D_b2 = tf.Variable(tf.zeros(shape=[1]))

# theta_D = [D_W1, D_W2, D_b1, D_b2]


z = tf.placeholder(tf.float32, shape=[None, z_dim])

# G_W1 = tf.Variable(xavier_init([z_dim, h_dim]))
# G_b1 = tf.Variable(tf.zeros(shape=[h_dim]))

# G_W2 = tf.Variable(xavier_init([h_dim, X_dim]))
# G_b2 = tf.Variable(tf.zeros(shape=[X_dim]))

# theta_G = [G_W1, G_W2, G_b1, G_b2]
# -

def get_sample_z(size):
    return np.random.uniform(-1., 1., size=size)


def generator(z):
    with tf.variable_scope('generator') as scope:
        G_h1 = tf.layers.dense(z, units=h_dim, kernel_initializer=tf.contrib.layers.xavier_initializer(),
                              activation=tf.nn.relu)
        G_out = tf.nn.sigmoid(tf.layers.dense(G_h1, units=X_dim, 
                                              kernel_initializer=tf.contrib.layers.xavier_initializer()))
        return G_out


def discriminator(x, reuse=False):
    with tf.variable_scope('discriminator', reuse=reuse) as scope:
        D_h1 = tf.layers.dense(x, units=h_dim, activation=tf.nn.relu, kernel_initializer=tf.contrib.layers.xavier_initializer())
        D_out = tf.layers.dense(D_h1, units=1, kernel_initializer=tf.contrib.layers.xavier_initializer())
        return D_out


generator_sample = generator(z)
discriminator_real = discriminator(X)
discriminator_fake = discriminator(generator_sample, True)

epsilon = tf.random_uniform([batch_size, 1], minval=0.0, maxval=1.0)
interpolated = epsilon * X + (1 - epsilon) * generator_sample
grads = tf.gradients(discriminator(interpolated, True), [interpolated])[0]
grad_norm = tf.sqrt(tf.reduce_sum((grads) ** 2, axis=1))
gp = lambda_grad * tf.reduce_mean((grad_norm - 1) ** 2)

d_loss = tf.reduce_mean(discriminator_fake) - tf.reduce_mean(discriminator_real) + gp
g_loss = -tf.reduce_mean(discriminator_fake)

# +
g_vars = [var for var in tf.trainable_variables() if var.name.startswith('generator')]
d_vars = [var for var in tf.trainable_variables() if var.name.startswith('discriminator')]

d_step = AdamirrorOptimizer(learning_rate=lr).minimize(d_loss, var_list=d_vars)
g_step = AdamirrorOptimizer(learning_rate=lr).minimize(g_loss, var_list=g_vars)
# -

print(g_vars)

print(d_vars)

sess = tf.InteractiveSession()

sess.run(tf.global_variables_initializer())

for i in range(100000):
    for _ in range(n_disc):
        x_batch, _ = mnist.train.next_batch(batch_size)

        _, d_loss_val = sess.run(
            [d_step, d_loss],
            feed_dict={X: x_batch, z: get_sample_z(size=(batch_size, z_dim))}
        )

    _, g_loss_val = sess.run(
        [g_step, g_loss],
        feed_dict={z: get_sample_z(size=(batch_size, z_dim))}
    )

    if i % 100 == 0:
        print('Iteration: {} - Discriminator Loss: {:.4}, Generator Loss: {:.4}'
              .format(i, -d_loss_val, g_loss_val))

        if i % 1000 == 0:
            samples = sess.run(generator_sample, feed_dict={z: get_sample_z(size=(16, z_dim))})

            fig = plot_images(samples)
            plot.show()
            i += 1
            plot.close(fig)



StarcoderdataJupyterScriptsDedupFiltered