Introduction

Generative Adversarial Network (GAN) is a class of generative deep learning models that is used  to learn the probability distribution of an underlying dataset and generates samples from the distribution. GAN consist of two deep learning networks

  • A Generator and

  • A Discriminator


The Generator network generates samples of data using randomly initialised network parameters(weights and bias). The task of discriminator network is to discriminate the real sample from generated samples

In this guide we provide the steps to train a GAN to generate hand-written digits using MNIST data.


Build the model architecture


1. Import the required packages

import razor
import numpy as np
from razor.flow import Pipeline, inputs, outputs

from rztdl.dl.components.layers import Input, Dense, Conv2DTranspose, Conv2D, MaxPool2D
from rztdl.dl.components.losses import MeanSquaredError
from rztdl.dl.components.operators import Add, Flatten, Reshape
from rztdl.dl.components.optimizers import Adam
from rztdl.dl.constants.string_constants import Padding
from rztdl.dl.groups import Group
from rztdl.dl.helpers.activations import Sigmoid
from rztdl.dl import RZTModel

2. Generator network

generator = Group("generator", scopes=["A"])
generator.add(Dense(name="dense_1", units=5, inputs="dense_input"))
generator.add(Dense(name="dense_2", units=128 * 14 * 14))
generator.add(Reshape(name="r1", target_shape=[14, 14, 128]))
generator.add(Conv2DTranspose(name="conv2D_transpose", filters=1, kernel_size=[2, 2], strides=[2, 2]))



3. Discriminator network

discriminator = Group("discriminator", scopes=['B'])
discriminator.add(Conv2D(name="conv1", kernel_size=2, filters=3, padding=Padding.SAME, inputs="conv1_in"))
discriminator.add(MaxPool2D(name="max_pool_1", pool_size=3, strides=2, padding=Padding.SAME))
discriminator.add(Flatten(name="flatten"))
discriminator.add(Dense(name="out", units=1, activation=Sigmoid()))

4. Input Layers

model = RZTModel("DC_GAN")
model.add(Input(name="noise", shape=[10]))
model.add(Input(name="images", shape=[28, 28, 1]))
model.add(Input(name="ones", shape=[1]))
model.add(Input(name="zeros", shape=[1]))


5. Connections

gen = generator("gen", inputs={"dense_input": 'noise'})
model.add(gen)
real_disc = discriminator("real_disc", inputs={"conv1_in": 'images'})
model.add(real_disc)
disc_fake = discriminator('disc_fake',inputs={"conv1_in": gen.get('conv2D_transpose')}, 
                                              shared="real_disc")
model.add(disc_fake)

model.add(MeanSquaredError(name="disc_mse_real", predictions=real_disc.get('out'), labels="ones"))

model.add(MeanSquaredError(name="disc_mse_fake", predictions=disc_fake.get('out'), labels="zeros"))

model.add(Add(name="discriminator_loss", inputs=['disc_mse_real', 'disc_mse_fake']))

model.add(MeanSquaredError(name="generator_loss", predictions=disc_fake.get('out'), labels='ones'))

model.add(Adam(name="generator_optimizer", inputs='generator_loss', scopes=['A']))

model.add(Adam(name="discriminator_optimizer", inputs='discriminator_loss', scopes=['B']))

model.plot()


Build the training pipeline


1. Data Reader Block

Define a block to read mnist data from project space

import razor.flow as rf
import pandas as pd
from razor.api import project_space_path
@rf.block
class DataReader():
    file_path:str
    chunk_size:int
    input_column_mapping:dict
    data: rf.SeriesOutput[dict]
    def run(self):
        chunks = pd.read_csv(project_space_path(self.file_path), chunksize=1)
        for df in chunks:
            k = {
            "noise": np.random.random((10)),
            #TODO remove the hard coding
            'images': df[self.input_column_mapping["Input_Image"]].values[0],
            'ones': np.ones((1)),
            'zeros': np.zeros((1))
            }
            self.data.put(k)

2. Import training, persist training and infer blocks


from   razor.marketplace.blocks.rzt.train_flow import DLTrainBlock
from   razor.marketplace.blocks.rzt.persist_flow import DLPersistBlock
from   razor.marketplace.blocks.rzt.infer_flow import DLInferBlock

3. Create the pipeline


trainer = DLTrainBlock("Gan_Trainer",
                     model =model,
                     train_data=data_reader.data,
                     test_data=data_reader.data,
                     epochs=1,
                     optimizers=["discriminator_optimizer"],
                     metrics=['discriminator_loss'],
                     batch_size=10,
                     save_path="gan",
                     log_frequency = 5
                      )
trainer.executor = rf.ContainerExecutor(cores=4, memory=4000)
persist_model = DLPersistBlock("DC_GAN_PERSIST",
                            model=model,
                            train_data=data_reader.data,
                            test_data=data_reader.data,
                            epochs=1,
                            optimizers=["generator_optimizer"],
                            metrics=['generator_loss'],
                            batch_size=10,
                            load_path=trainer.saved_model,
                            save_path="gan",
                            log_frequency=5)
persist_model.executor = rf.ContainerExecutor(cores=4, memory=4000)

infer_model = DLInferBlock(name="gan_infer",
                        model=model,
                        batch_size=10,
                        data=data_reader.data,
                        load_path=persist_model.saved_model,
                        layers=['generator_loss']
                        )
infer_model.executor = rf.ContainerExecutor(cores=4, memory=4000)
pipeline = Pipeline(targets=[infer_model])
pipeline.show()



4. Executing the pipeline in engine

Replace ```<Engine Name>``` with the name of engine in the following code


razor.api.engines(<Engine Name>).execute(pipeline)



We can view the status of the pipeline in the Engine through IDE