import tensorflow as tf
from tensorflow.keras import Model, layers, backend
from tensorflow.keras.constraints import Constraint
from losses import disc_hinge, disc_loss, gen_loss, gen_hinge
from diff_augment import diff_augment
from tensorflow_addons.layers import SpectralNormalization
tf.random.set_seed(45)
# np.random.seed(45)
class Generator(Model):
def __init__(self, n_class=10, res=128):
super(Generator, self).__init__()
# filters = [ 1024, 512, 256, 128, 64, 32]#, 32, 16]
# strides = [ 4, 2, 2, 2, 2, 2]#, 2, 2]
filters = [ 1024, 512, 256, 128, 64, 32]#, 16]
strides = [ 4, 2, 2, 2, 2, 2]#, 2]
self.cnn_depth = len(filters)
# For discrete condition we are using Embedding
self.cond_embedding = layers.Embedding(input_dim=n_class, output_dim=50)
self.cond_flat = layers.Flatten()
self.cond_dense = layers.Dense(units=(8 * 8 * 1))
self.cond_reshape = layers.Reshape(target_shape=(64,))
# Hyperparameter:
# If only conv : mean=0.0, var=0.02
# If using bnorm: mean=1.0, var=0.02
self.conv = [SpectralNormalization(layers.Conv2DTranspose(\
filters=filters[idx], kernel_size=3,\
strides=strides[idx], padding='same',\
kernel_initializer=tf.keras.initializers.TruncatedNormal(mean=0.0, stddev=0.02),\
use_bias=False))\
for idx in range(self.cnn_depth)]
self.act = [layers.LeakyReLU() for idx in range(self.cnn_depth)]
self.bnorm = [layers.BatchNormalization() for idx in range(self.cnn_depth)]
self.last_conv = SpectralNormalization(layers.Conv2D(filters=3, kernel_size=3,\
strides=1, padding='same',\
activation='tanh',\
kernel_initializer=tf.keras.initializers.TruncatedNormal(mean=0.0, stddev=0.02),\
use_bias=False))
@tf.function
def call(self, X):
# C = self.cond_reshape( self.cond_dense( self.cond_flat( self.cond_embedding( C ) ) ) )
# X = tf.concat([C, X], axis=-1)
X = tf.expand_dims(tf.expand_dims(X, axis=1), axis=1)
X = self.act[0]( self.conv[0]( X ) )
for idx in range(1, self.cnn_depth):
X = self.act[idx]( self.bnorm[idx]( self.conv[idx]( X ) ) )
# X = self.bnorm[idx]( self.act[idx]( self.conv[idx]( X ) ) )
# X = self.act[idx]( self.conv[idx]( X ) )
X = self.last_conv(X)
return X
class Discriminator(Model):
def __init__(self, n_class=10, res=128):
super(Discriminator, self).__init__()
# filters = [32, 64, 128, 256, 256, 512, 512, 1]
# strides = [ 2, 2, 2, 2, 2, 2, 1, 1]
# filters = [ 64, 128, 256, 512, 1024, 1]
# strides = [ 2, 2, 2, 2, 1, 1]
filters = [ 64, 128, 256, 512, 1024, 1]
strides = [ 2, 2, 2, 2, 1, 1]
self.cnn_depth = len(filters)
# For discrete condition we are using Embedding
self.cond_embedding = layers.Embedding(input_dim=n_class, output_dim=50)
self.cond_flat = layers.Flatten()
self.cond_dense = layers.Dense(units=(res * res * 1))
self.cond_reshape = layers.Reshape(target_shape=(res, res, 1))
self.cnn_conv = [layers.Conv2D(filters=filters[i], kernel_size=3,\
strides=strides[i], padding='same',\
kernel_initializer=tf.keras.initializers.TruncatedNormal(mean=0.0, stddev=0.02),\
use_bias=False)\
for i in range(self.cnn_depth)]
self.cnn_bnorm = [layers.BatchNormalization() for _ in range(self.cnn_depth)]
self.cnn_act = [layers.LeakyReLU(alpha=0.2) for _ in range(self.cnn_depth)]
# self.final_act = layers.Activation('sigmoid')
self.flat = layers.Flatten()
self.disc_out = layers.Dense(units=1)
# self.autoenc = Autoencoder()
@tf.function
def call(self, x, C):
#x = self.cnn_merge( x )
#x = self.cnn_exp( x )
# mem_bank = []
# C = self.cond_reshape( self.cond_dense( self.cond_flat( self.cond_embedding( C ) ) ) )
C = tf.expand_dims( tf.expand_dims(C, axis=1), axis=1)
C = tf.tile(C, [1, x.shape[1], x.shape[2], 1])
x = tf.concat([x, C], axis=-1)
for layer_no in range(self.cnn_depth):
# print(x.shape)
x = self.cnn_act[layer_no]( self.cnn_bnorm[layer_no]( self.cnn_conv[layer_no]( x ) ) )
# x = self.cnn_bnorm[layer_no]( self.cnn_act[layer_no]( self.cnn_conv[layer_no]( x ) ) )
# x = self.cnn_act[layer_no]( self.cnn_conv[layer_no]( x ) )
# if layer_no == 0:
# mem_bank.append( x )
# if layer_no == 1:
# mem_bank.append( x )
# x = self.cnn_act[layer_no]( self.cnn_conv[layer_no]( x ) )
# reconst_x = self.autoenc( x )
# condition = tf.expand_dims(tf.expand_dims(condition, axis=1), axis=1)
# condition = tf.tile(condition, [1, x.shape[1], x.shape[1], 1])
# x = tf.concat([x, condition], axis=-1)
# x = self.cnn_act[layer_no+1]( self.cnn_bnorm[layer_no+1]( self.cnn_conv[layer_no+1]( x ) ) )
# x = self.cnn_bnorm[layer_no+1]( self.cnn_act[layer_no+1]( self.cnn_conv[layer_no+1]( x ) ) )
# x = self.cnn_act[layer_no+1]( self.cnn_conv[layer_no+1]( x ) )
# reconst_x = self.autoenc( x )
reconst_x = None
# x = self.cnn_act[layer_no+2]( self.cnn_bnorm[layer_no+2]( self.cnn_conv[layer_no+2]( x ) ) )
# reconst_x = self.autoenc( x, mem_bank )
# x = self.final_act( x )
# x = self.out( self.flat( x ) )
x = self.disc_out( self.flat( x ) )
return x, reconst_x
class DCGAN(Model):
def __init__(self):
super(DCGAN, self).__init__()
self.gen = Generator()
self.disc = Discriminator()
@tf.function
def dist_train_step(mirrored_strategy, model, model_gopt, model_copt, X, C, latent_dim=96, batch_size=64):
diff_augment_policies = "color,translation"
noise_vector = tf.random.uniform(shape=(batch_size, latent_dim), minval=-1, maxval=1)
noise_vector_2 = tf.random.uniform(shape=(batch_size, latent_dim), minval=-1, maxval=1)
noise_vector = tf.concat([noise_vector, C], axis=-1)
noise_vector_2 = tf.concat([noise_vector_2, C], axis=-1)
# @tf.function
def train_step_disc(model, model_gopt, model_copt, X, C, latent_dim=96, batch_size=64):
with tf.GradientTape() as ctape:
# noise_vector = tf.random.uniform(shape=(batch_size, latent_dim), minval=-1, maxval=1)
# noise_vector = tf.random.uniform(shape=(batch_size, latent_dim), minval=-1, maxval=1)
# noise_vector = tf.random.normal(shape=(batch_size, latent_dim))
fake_img = model.gen(noise_vector, training=False)
X_aug = diff_augment(X, policy=diff_augment_policies)
fake_img = diff_augment(fake_img, policy=diff_augment_policies)
D_real, X_recon = model.disc(X_aug, C, training=True)
D_fake, _ = model.disc(fake_img, C, training=True)
# c_loss = disc_loss(D_real, D_fake) +\
# tf.reduce_mean(tf.keras.losses.MeanAbsoluteError(reduction=tf.keras.losses.Reduction.NONE)(X_aug, X_recon))
# c_loss = disc_hinge(D_real, D_fake) +\
# tf.reduce_mean(tf.keras.losses.MeanAbsoluteError(reduction=tf.keras.losses.Reduction.NONE)(X_aug, X_recon))
c_loss = disc_hinge(D_real, D_fake)
variables = model.disc.trainable_variables
gradients = ctape.gradient(c_loss, variables)
model_copt.apply_gradients(zip(gradients, variables))
return c_loss
# @tf.function
def train_step_gen(model, model_gopt, model_copt, X, C, latent_dim=96, batch_size=64):
with tf.GradientTape() as gtape:
# noise_vector = tf.random.uniform(shape=(batch_size, latent_dim), minval=-1, maxval=1)
# noise_vector = tf.random.normal(shape=(batch_size, latent_dim))
fake_img_o = model.gen(noise_vector, training=True)
fake_img_2_o = model.gen(noise_vector_2, training=True)
#D_fake = model.disc(fake_img, H_hat, training=False)
fake_img = diff_augment(fake_img_o, policy=diff_augment_policies)
fake_img_2 = diff_augment(fake_img_2_o, policy=diff_augment_policies)
D_fake, _ = model.disc(fake_img, C, training=False)
D_fake_2, _ = model.disc(fake_img_2, C, training=False)
# g_loss = gen_loss(D_fake)
g_loss = gen_hinge(D_fake) + gen_hinge(D_fake_2)
mode_loss = tf.divide(tf.reduce_mean(tf.abs(tf.subtract(fake_img_2_o, fake_img_o))),\
tf.reduce_mean(tf.abs(tf.subtract(noise_vector_2, noise_vector)))
)
mode_loss = tf.divide(1.0, mode_loss + 1e-5)
g_loss = g_loss + 1.0 * mode_loss
variables = model.gen.trainable_variables #+ model.gcn.trainable_variables
gradients = gtape.gradient(g_loss, variables)
model_gopt.apply_gradients(zip(gradients, variables))
return g_loss
per_replica_loss_disc = mirrored_strategy.run(train_step_disc, args=(model, model_gopt, model_copt, X, C, latent_dim, batch_size,))
per_replica_loss_gen = mirrored_strategy.run(train_step_gen, args=(model, model_gopt, model_copt, X, C, latent_dim, batch_size,))
# print(per_replica_loss_disc)
# print(mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss_disc, axis=0).numpy())
discriminator_loss = mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss_disc, axis=None)
generator_loss = mirrored_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss_gen, axis=None)
return generator_loss, discriminator_loss
Datasets
Models
