Spaces:
Sleeping
Sleeping
File size: 6,395 Bytes
5b8ac50 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
from tensorflow.keras import Sequential, layers
from tensorflow.keras.preprocessing.image import array_to_img
# Declerating the configuration params
LATENT_DIM = 100
WEIGHT_INIT = keras.initializers.RandomNormal(mean=0.0, stddev=0.02)
CHANNELS = 3
input_shape = (64, 64, 3)
alpha = 0.2
def get_generator():
# Model architecture
generator = Sequential(name='generator')
# 1d random noise
generator.add(layers.Dense(8 * 8 * 512, input_dim=LATENT_DIM))
# model.add(layers.BatchNormalization())
generator.add(layers.ReLU())
# convert 1d to 3d
generator.add(layers.Reshape((8, 8, 512)))
# upsample to 16x16
generator.add(layers.Conv2DTranspose(256, (4, 4), strides=(2, 2), padding='same', kernel_initializer=WEIGHT_INIT))
# model.add(layers.BatchNormalization())
generator.add(layers.ReLU())
# upsample to 32x32
generator.add(layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding='same', kernel_initializer=WEIGHT_INIT))
# model.add(layers.BatchNormalization())
generator.add(layers.ReLU())
# upsample to 64x64
generator.add(layers.Conv2DTranspose(64, (4, 4), strides=(2, 2), padding='same', kernel_initializer=WEIGHT_INIT))
# model.add(layers.BatchNormalization())
generator.add(layers.ReLU())
# Output layer
generator.add(layers.Conv2D(CHANNELS, (4, 4), padding='same', activation='tanh'))
return generator
def get_discriminator():
# Creating the architecture
discriminator = Sequential(name='discriminator')
# conv layer-1
discriminator.add(layers.Conv2D(64, (4, 4), strides=(2, 2), padding='same', input_shape=input_shape))
discriminator.add(layers.BatchNormalization())
discriminator.add(layers.LeakyReLU(alpha=alpha))
# conv layer-2
discriminator.add(layers.Conv2D(128, (4, 4), strides=(2, 2), padding='same', input_shape=input_shape))
discriminator.add(layers.BatchNormalization())
discriminator.add(layers.LeakyReLU(alpha=alpha))
# conv layer-3
discriminator.add(layers.Conv2D(128, (4, 4), strides=(2, 2), padding='same', input_shape=input_shape))
discriminator.add(layers.BatchNormalization())
discriminator.add(layers.LeakyReLU(alpha=alpha))
# Fully Connected Classifier layer
discriminator.add(layers.Flatten())
discriminator.add(layers.Dropout(0.3))
# output class
discriminator.add(layers.Dense(1, activation='sigmoid'))
return discriminator
class DCGAN(keras.Model):
def __init__(self, generator, discriminator, latent_dim):
super().__init__()
self.generator = generator
self.discriminator = discriminator
self.latent_dim = latent_dim
self.g_loss_metric = keras.metrics.Mean(name='g_loss')
self.d_loss_metric = keras.metrics.Mean(name='d_loss')
@property
def metrics(self):
return [self.g_loss_metric, self.d_loss_metric]
def compile(self, g_optimizer, d_optimizer, loss_fn):
super(DCGAN, self).compile()
self.g_optimizer = g_optimizer
self.d_optimizer = d_optimizer
self.loss_fn = loss_fn
def train_step(self, real_images):
# generate random noise
batch_size = tf.shape(real_images)[0]
random_noise = tf.random.normal(shape=(batch_size, self.latent_dim))
# train the discriminator with real (1) and fake (0) images
with tf.GradientTape() as tape:
# discriminator (pred then loss)
pred_real = self.discriminator(real_images, training=True)
real_labels = tf.ones((batch_size, 1))
real_labels += 0.05 * tf.random.uniform(tf.shape(real_labels))
d_loss_real = self.loss_fn(real_labels, pred_real)
# generator (generate then implemen discriminator steps again)
fake_images = self.generator(random_noise)
pred_fake = self.discriminator(fake_images, training=True)
fake_labels = tf.zeros((batch_size, 1))
d_loss_fake = self.loss_fn(fake_labels, pred_fake)
# total discriminator loss
d_loss = (d_loss_real + d_loss_fake) / 2
# compute discriminator gradients then update the gradients
gradients = tape.gradient(d_loss, self.discriminator.trainable_variables)
self.d_optimizer.apply_gradients(zip(gradients, self.discriminator.trainable_variables))
# train the generator model
labels = tf.ones((batch_size, 1))
with tf.GradientTape() as tape:
# generate then implement discriminator
fake_images = self.generator(random_noise, training=True)
pred_fake = self.discriminator(fake_images, training=True)
g_loss = self.loss_fn(labels, pred_fake)
# compute gradients then update the gradients
gradients = tape.gradient(g_loss, self.generator.trainable_variables)
self.g_optimizer.apply_gradients(zip(gradients, self.generator.trainable_variables))
# update states for both models
self.d_loss_metric.update_state(d_loss)
self.g_loss_metric.update_state(g_loss)
return {'d_loss': self.d_loss_metric.result(), 'g_loss': self.g_loss_metric.result()}
# Custom Callback to display image per every train step
class DCGANMonitor(keras.callbacks.Callback):
def __init__(self, num_imgs=25, latent_dim=100):
self.num_imgs = num_imgs
self.latent_dim = latent_dim
self.noise = tf.random.normal([25, latent_dim])
def on_epoch_end(self, epoch, logs=None):
# generate the image from noise
g_img = self.model.generator(self.noise)
# denormalize the image
g_img = (g_img * 127.5) + 127.5
g_img.numpy()
# plot the image per step
fig = plt.figure(figsize=(8, 8))
for i in range(self.num_imgs):
plt.subplot(5, 5, i+1)
img = array_to_img(g_img[i])
plt.imshow(img)
plt.axis('off')
# plt.savefig('epoch_{:03d}.png'.format(epoch))
plt.show()
def on_train_end(self, logs=None):
self.model.generator.save('generator.h5') |