I am trying to implement a WGAN loss function with Gradient Penalty on TPU. After training, the result is not what I expected it to be.
Below is the graph
So,What I expected:
-
I expected a continuous decrease in both Generator and Discriminator loss.
-
The values should have been under a certain limit.
My code for the Generator and Critic(Discriminator) Loss:
class CriticLoss(object):
""" Criric Loss
Args:
discriminator:Discriminator model
Dx: Output of the real images from discriminator
Dx_hat: Output of the generated images(fake) from discriminator
x_interpolated:combined fake and real images
"""
def __init__(self, gp_lambda=10):
self.gp_lambda = gp_lambda
def __call__(self,discriminator, Dx, Dx_hat,x_interpolated):
#orgnal critic loss
d_loss = tf.reduce_mean(Dx_hat) - tf.reduce_mean(Dx)
#calculate gradinet penalty
with tf.GradientTape() as tape:
tape.watch(x_intepolated)
dx_inter = discriminator(x_interpolated, training=True)
gradients=tape.gradient(dx_inter, [x_interpolated])[0]
grad_l2 = tf.sqrt(tf.reduce_sum(tf.square(gradients), axis=[1, 2, 3]))
grad_penalty = tf.reduce_mean(tf.square(grad_l2 - 1.0))
#final discriminator loss
d_loss += self.gp_lambda * grad_penalty
return d_loss
#Generator loss
class GeneratorLoss(object):
""" Generator Loss """
def __call__(self,Dx_hat):
return tf.reduce_mean(-Dx_hat)
Since, I already checked my DCGAN model with CrossEntropy Loss and It works perfectly fine.So my model is not in fault here.It could be the fact that how TPU distribution strategy works and the loss functions calculated in the individual TPU device might not adding up to provide suitable values.
Also, I should point out that the loss values in the graph are calculated in the following way.
gen_loss.update_state(g_loss * tpu_strategy.num_replicas_in_sync)
disc_loss.update_state(d_loss * tpu_strategy.num_replicas_in_sync )
where gen_loss and disc_loss are defined as tf.keras.metrics.Mean() inside tpu_strategy.scope() while g_loss and d_loss are the output values from the GeneratorLoss and CriticLoss repsectively in the step_fn