update: codes

kozistr · Oct 12, 2020 · 4e27ef0 · 4e27ef0
1 parent 620b931
commit 4e27ef0
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 10 deletions.
diff --git a/awesome_gans/data.py b/awesome_gans/data.py
@@ -36,7 +36,7 @@ def load_dataset(self, use_label: bool = False):
         ds = tfds.load(name=self.dataset, split='train', as_supervised=use_label, shuffle_files=True)
         ds = ds.map(lambda x: self.preprocess_image(x['image']), tf.data.experimental.AUTOTUNE)
         ds = ds.cache()
-        ds = ds.shuffle(50000)
+        ds = ds.shuffle(self.bs * 16)
         ds = ds.batch(self.bs, drop_remainder=True)
         ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
         return ds
diff --git a/awesome_gans/losses.py b/awesome_gans/losses.py
@@ -1,6 +1,17 @@
 import tensorflow as tf
 
 
+@tf.function
+def discriminator_wgan_loss(real: tf.Tensor, fake: tf.Tensor):
+    return tf.reduce_mean(fake) - tf.reduce_mean(real)
+
+
+@tf.function
+def generator_wgan_loss(fake: tf.Tensor):
+    return -tf.reduce_mean(fake)
+
+
+@tf.function
 def discriminator_loss(loss_func: str, real: tf.Tensor, fake: tf.Tensor, use_ra: bool = False):
     real_loss: float = 0.0
     fake_loss: float = 0.0
@@ -31,6 +42,7 @@ def discriminator_loss(loss_func: str, real: tf.Tensor, fake: tf.Tensor, use_ra:
     return loss
 
 
+@tf.function
 def generator_loss(loss_func: str, real: tf.Tensor, fake: tf.Tensor, use_ra: bool = False):
     fake_loss: float = 0.0
     real_loss: float = 0.0

diff --git a/awesome_gans/wgan/model.py b/awesome_gans/wgan/model.py
@@ -16,7 +16,7 @@
 from tensorflow.keras.models import Model
 from tqdm import tqdm
 
-from awesome_gans.losses import discriminator_loss, generator_loss
+from awesome_gans.losses import discriminator_loss, generator_loss, discriminator_wgan_loss, generator_wgan_loss
 from awesome_gans.optimizers import build_optimizer
 from awesome_gans.utils import merge_images, save_image
 
@@ -103,15 +103,15 @@ def train_discriminator(self, x: tf.Tensor):
             d_fake = self.discriminator(x_fake, training=True)
             d_real = self.discriminator(x, training=True)
 
-            d_loss = discriminator_loss(self.d_loss, d_real, d_fake)
+            d_loss = discriminator_wgan_loss(d_real, d_fake)
 
-            gradient = gt.gradient(d_loss, self.discriminator.trainable_variables)
-            self.d_opt.apply_gradients(zip(gradient, self.discriminator.trainable_variables))
+            gradients = gt.gradient(d_loss, self.discriminator.trainable_variables)
+            self.d_opt.apply_gradients(zip(gradients, self.discriminator.trainable_variables))
 
             for var in self.discriminator.trainable_variables:
                 var.assign(tf.clip_by_value(var, -self.grad_clip, self.grad_clip))
 
-        return d_loss
+            return d_loss
 
     @tf.function
     def train_generator(self):
@@ -120,12 +120,12 @@ def train_generator(self):
             x_fake = self.generator(z, training=True)
             d_fake = self.discriminator(x_fake, training=True)
 
-            g_loss = generator_loss(self.g_loss, d_fake, d_fake)
+            g_loss = generator_wgan_loss(d_fake)
 
-            gradient = gt.gradient(g_loss, self.generator.trainable_variables)
-            self.g_opt.apply_gradients(zip(gradient, self.generator.trainable_variables))
+            gradients = gt.gradient(g_loss, self.generator.trainable_variables)
+            self.g_opt.apply_gradients(zip(gradients, self.generator.trainable_variables))
 
-        return g_loss
+            return g_loss
 
     def load(self) -> int:
         return 0