from: https://medium.com/aubergine-solutions/implementing-gan-using-tensorflow-2-0-9d03e29692cf

In this article, we’ll understand how to implement Generative Adversarial Networks in Tensorflow 2.0 using the MNIST dataset of digits. Our model should be able to accurately create images similar to MNIST dataset images. Image from MNIST dataset looks like this:

Prerequisites

Before moving ahead please make sure you are updated with the following requirements:

Basic understanding of Tensorflow 2.0
Basic understanding of GANs
Tensorflow 2.0 library with Python 3.6

Let’s jump into code directly

For implementing the network, we need to follow the following steps.

Step 1: Import Libraries

We will need to import these libraries.

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
import os
from tensorflow.keras import layers
import time

from IPython import display

Step 2: Load and Prepare data

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(60000).batch(256)

In this piece of code, we imported data and reshaped it. Images are represented in the values between 0 to 255. So these image vectors are converted into values between [-1, 1]. Then these vectors are shuffled and converted to batches of 256 images.

Step 3: Create Models

The generator model will take vector of 100 pixels as input and convert that vector into an image of 26 * 26 using the Conv2DTranspose layer. As we are using convolutions, this model will be a DCGAN i.e. Deep Convolutional Generative Adversarial Network.

BatchNormalization layer is used for normalizing the generated image to reduce noise that image. The activation function of “ReLU” is added as another layer in the model.

def create_generator():
    model = tf.keras.Sequential()
    
    # creating Dense layer with units 7*7*256(batch_size) and input_shape of (100,)
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 256)))

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))

    return model

Discriminator will do the exact opposite of a generator and convert image into scalar probabilities of whether the image is fake or real. It will use the Conv2D layer for this purpose. Along with the convolutional layer, it will have layer of an activation function “ReLU” and the Dropout layer.

def create_discriminator():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[28, 28, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(1))

    return model

After both the models are created, we’ll move to the next steps.

Step 4: Define Loss and Optimizers

Here, we’re using BinaryCrossentropy from tf.keras.losses API. Using this cross_entropy loss, we’ll make loss functions for discriminator and generator.

cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

def D_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss
  
def G_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

训练判别器，（或固定生成器，）期望判别器识别真实图片输出为1，生成图片输出为0

训练生成器，（或固定判别器，）期望判别器识别生成图片输出为1 （欺骗）

Now, we’ll create optimizers for generator and discriminator using Adam optimizer from tf.keras.optimizers API with a learning rate of 1e-4.

generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)

Step 5: Define Training functions

Initially, the generator will generate an image with only random pixels and as time passes, it will make images that are almost similar to the training images. Image from the generator before training looks like the following:

noise_dim = 100
num_of_generated_examples = 16

seed = tf.random.normal([num_of_generated_examples, noise_dim])

First, we need to create noise data for passing into the generator model. Then create a function to learn through iterations using the concept of Eager Execution. Please refer to this for understanding what is Eager Execution as it is out of scope for this tutorial.

generator = create_generator()
discriminator = create_discriminator()

@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      generated_images = generator(noise, training=True)

      real_output = discriminator(images, training=True)
      fake_output = discriminator(generated_images, training=True)

      gen_loss = G_loss(fake_output)
      disc_loss = D_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

GradientTape above is used for creating gradients as we go further and these gradients are applied to optimizers for decreasing the loss function.

def train_GAN(dataset, epochs):
  for epoch in range(epochs):
    start = time.time()

    for image_batch in dataset:
      train_step(image_batch)
      
     print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

This is the final function to call the train_step function for each batch in the dataset.

Finally, we’ll call our train_GAN function to start training on the image dataset with 500 epochs.

train_GAN(train_dataset, 500)

After going through all the epochs, the generated image looks like this:

Seeing the results

Let’s look at the results in the form of GIF.

It could be seen how our neural networks learn to write integers like a child. We can use the same algorithm for multiple image generators because we only need to change the training images and train again.

What’s next?

GANs could be used for many more domains other than images with the same concept of Generator and Discriminator. Check out this article for learning what are different types of GANs.

Implementing GAN using Tensorflow 2.0