Add artificial balancing

1 files changed, 7 insertions, 4 deletions

diff --git a/ncdcgan.py b/ncdcgan.py
index 97b137b..65c5862 100755
--- a/ncdcgan.py
+++ b/ncdcgan.py
@@ -148,7 +148,7 @@ class nCDCGAN():
         model.summary()
         return model
 
-    def train(self, epochs, batch_size=128, sample_interval=50, graph=False, smooth_real=1, smooth_fake=0):
+    def train(self, epochs, batch_size=128, sample_interval=50, graph=False, smooth_real=1, smooth_fake=0, gdbal = 1):
 
         # Load the dataset
         (X_train, y_train), (_, _) = mnist.load_data()
@@ -181,9 +181,12 @@ class nCDCGAN():
             gen_imgs = self.generator.predict([noise, labels])
 
             # Train the discriminator
-            d_loss_real = self.discriminator.train_on_batch([imgs, labels], valid*smooth_real)
-            d_loss_fake = self.discriminator.train_on_batch([gen_imgs, labels], valid*smooth_fake)
-            d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
+            if epoch % gdbal == 0:
+                d_loss_real = self.discriminator.train_on_batch([imgs, labels], valid*smooth_real)
+                d_loss_fake = self.discriminator.train_on_batch([gen_imgs, labels], valid*smooth_fake)
+                d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
+            else:
+                dloss = 0
 
             # ---------------------
             #  Train Generator