Generative Adversarial Network (GAN) using RZT aiOS SDK
More articles in Tips, Tricks, and More
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().execute(pipeline)
