Training Rate Scaling

Author: Aggrathon

generator_gan.py

@@ -91,7 +91,7 @@ def setup_network(self):
         dis_var = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='discriminator')
         self.generator_solver, self.discriminator_solver, self.scale = \
             gan_optimizer('train', gen_var, dis_var, gen_logit, image_logit, self._y_offset, 1-self._y_offset,
-                          *self.learning_rate, global_step=self.iterations, summary=self.log)
+                          *self.learning_rate, learning_rate_pivot=10000, global_step=self.iterations, summary=self.log)
 
 
     def random_input(self):
@@ -190,10 +190,14 @@ def train(self, batches=100000, print_interval=1):
                 if timer() - last_save > 1800:
                     saver.save(session, os.path.join(self.directory, self.name))
                     last_save = timer()
+            print("Saving the network")
+            saver.save(session, os.path.join(self.directory, self.name))
+            if self.log:
+                logger.close()
+            session.close()
         except KeyboardInterrupt:
             print()
-            print("Stopping the training", end='')
-        finally:
+            print("Stopping the training")
             saver.save(session, os.path.join(self.directory, self.name))
             if self.log:
                 logger.close()

network.py

@@ -91,9 +91,14 @@ def batch_optimizer(name, variables, positive_tensors=None, positive_value=1, po
         return solver
 
 def gan_optimizer(name, gen_vars, dis_vars, fake_tensor, real_tensor, false_val=0, real_val=1,
-                  learning_rate=0.001, learning_momentum=0.9, learning_momentum2=0.99, global_step=None, summary=True):
+                  learning_rate=0.001, learning_momentum=0.9, learning_momentum2=0.99,
+                  learning_rate_pivot=0, global_step=None, summary=True):
     """Create an optimizer for a GAN"""
     with openif_scope(name):
+        #learning rate scaling
+        if learning_rate_pivot > 0:
+            scaler = tf.sqrt(tf.div(tf.to_float(global_step), float(learning_rate_pivot))+1)
+            learning_rate = tf.div(learning_rate, scaler)
         #generator
         with tf.variable_scope('generator'):
             gen_labels = tf.fill(tf.shape(fake_tensor), real_val)

train.py

@@ -5,7 +5,7 @@
 
 CONFIG = {
     'colors': 3,
-    'batch_size': 128,
+    'batch_size': 192,
     'generator_base_width': 32,
     'image_size': 64,
     'discriminator_convolutions': 5,