def train_step(self, data):
input_image, target = data
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
gen_output = self.generator(input_image, training=True)
disc_real_output = self.discriminator([input_image, target],
training=True)
disc_generated_output = self.discriminator(
[input_image, gen_output], training=True)
gen_total_loss, gen_gan_loss, gen_l1_loss = self.generator_loss(
disc_generated_output, gen_output, target)
disc_loss = self.discriminator_loss(disc_real_output,
disc_generated_output)
generator_gradients = gen_tape.gradient(
gen_total_loss, self.generator.trainable_variables)
discriminator_gradients = disc_tape.gradient(
disc_loss, self.discriminator.trainable_variables)
self.generator_optimizer.apply_gradients(
zip(generator_gradients, self.generator.trainable_variables))
self.discriminator_optimizer.apply_gradients(
zip(discriminator_gradients,
self.discriminator.trainable_variables))
return {
"d_loss": disc_loss,
"g_total": gen_total_loss,
'g_gan': gen_gan_loss,
'g_l1': gen_l1_loss,
'dice': dice(target, gen_output)
}
def test_step(self, data):
input_image, target = data
pred = self.generator(input_image, training=False)
print(target)
return {'dice': dice(target, pred)}