def run_gan():
(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()
print(train_images.shape)
train_images = train_images.reshape(train_images.shape[0], 28, 28,
1).astype('float32')
train_images = (train_images - 127.5) / 127.5 # Normalize images to [-1,1]
print(train_images.shape)
# Batch and shuffle the data
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(
BUFFER_SIZE).batch(BATCH_SIZE)
gan = DCGAN(gen_lr, disc_lr, batch_size=BATCH_SIZE, noise_dim=NOISE_DIM)
gan.create_generator()
gan.create_discriminator()
# Test generator
random_noise = tf.random.normal([1, NOISE_DIM])
generated_image = gan.generator(random_noise)
#plt.imshow(generated_image[0,:,:,0],cmap='gray')
#plt.show()
# Test Discriminator
prob = gan.discriminator(generated_image)
print("Probability of image being real: {}".format(sigmoid(prob)))
gan.set_noise_seed(num_examples_to_generate)
gan.set_checkpoint(path=save_ckpt_path)
gen_loss_array, disc_loss_array = gan.train(train_dataset, epochs=EPOCHS)
# Plot Discriminator Loss
plt.plot(range(EPOCHS), gen_loss_array)
plt.plot(range(EPOCHS), disc_loss_array)
plt.show()
