Update train.py

Author: zsdonghao

train.py

@@ -16,41 +16,44 @@ def train():
     G.train()
     D.train()
 
-    d_optimizer = tf.optimizers.Adam(flags.learning_rate, beta_1=flags.beta1)
-    g_optimizer = tf.optimizers.Adam(flags.learning_rate, beta_1=flags.beta1)
+    d_optimizer = tf.optimizers.Adam(flags.lr, beta_1=flags.beta1)
+    g_optimizer = tf.optimizers.Adam(flags.lr, beta_1=flags.beta1)
 
     n_step_epoch = int(len(images_path) // flags.batch_size)
-
-    for step, batch_images in enumerate(images):
-        step_time = time.time()
-        with tf.GradientTape(persistent=True) as tape:
-            # z = tf.distributions.Normal(0., 1.).sample([flags.batch_size, flags.z_dim]) #tf.placeholder(tf.float32, [None, z_dim], name='z_noise')
-            z = np.random.normal(loc=0.0, scale=1.0, size=[flags.batch_size, flags.z_dim]).astype(np.float32)
-            d_logits = D(G(z))
-            d2_logits = D(batch_images)
-            # discriminator: real images are labelled as 1
-            d_loss_real = tl.cost.sigmoid_cross_entropy(d2_logits, tf.ones_like(d2_logits), name='dreal')
-            # discriminator: images from generator (fake) are labelled as 0
-            d_loss_fake = tl.cost.sigmoid_cross_entropy(d_logits, tf.zeros_like(d_logits), name='dfake')
-            # combined loss for updating discriminator
-            d_loss = d_loss_real + d_loss_fake
-            # generator: try to fool discriminator to output 1
-            g_loss = tl.cost.sigmoid_cross_entropy(d_logits, tf.ones_like(d_logits), name='gfake')
-
-        grad = tape.gradient(g_loss, G.trainable_weights)
-        g_optimizer.apply_gradients(zip(grad, G.trainable_weights))
-        grad = tape.gradient(d_loss, D.trainable_weights)
-        d_optimizer.apply_gradients(zip(grad, D.trainable_weights))
-        del tape
-
-        print("Epoch: [{}/{}] [{}/{}] took: {:3f}, d_loss: {:5f}, g_loss: {:5f}".format(step//n_step_epoch, flags.n_epoch, step, n_step_epoch, time.time()-step_time, d_loss, g_loss))
-        if np.mod(step, flags.save_step) == 0:
+    
+    for epoch in range(flags.n_epoch):
+        for step, batch_images in enumerate(images):
+            step_time = time.time()
+            with tf.GradientTape(persistent=True) as tape:
+                # z = tf.distributions.Normal(0., 1.).sample([flags.batch_size, flags.z_dim]) #tf.placeholder(tf.float32, [None, z_dim], name='z_noise')
+                z = np.random.normal(loc=0.0, scale=1.0, size=[flags.batch_size, flags.z_dim]).astype(np.float32)
+                d_logits = D(G(z))
+                d2_logits = D(batch_images)
+                # discriminator: real images are labelled as 1
+                d_loss_real = tl.cost.sigmoid_cross_entropy(d2_logits, tf.ones_like(d2_logits), name='dreal')
+                # discriminator: images from generator (fake) are labelled as 0
+                d_loss_fake = tl.cost.sigmoid_cross_entropy(d_logits, tf.zeros_like(d_logits), name='dfake')
+                # combined loss for updating discriminator
+                d_loss = d_loss_real + d_loss_fake
+                # generator: try to fool discriminator to output 1
+                g_loss = tl.cost.sigmoid_cross_entropy(d_logits, tf.ones_like(d_logits), name='gfake')
+
+            grad = tape.gradient(g_loss, G.trainable_weights)
+            g_optimizer.apply_gradients(zip(grad, G.trainable_weights))
+            grad = tape.gradient(d_loss, D.trainable_weights)
+            d_optimizer.apply_gradients(zip(grad, D.trainable_weights))
+            del tape
+
+            print("Epoch: [{}/{}] [{}/{}] took: {:3f}, d_loss: {:5f}, g_loss: {:5f}".format(epoch, \
+                  flags.n_epoch, step, n_step_epoch, time.time()-step_time, d_loss, g_loss))
+        
+        if np.mod(epoch, flags.save_every_epoch) == 0:
             G.save_weights('{}/G.npz'.format(flags.checkpoint_dir), format='npz')
             D.save_weights('{}/D.npz'.format(flags.checkpoint_dir), format='npz')
             G.eval()
             result = G(z)
             G.train()
-            tl.visualize.save_images(result.numpy(), [num_tiles, num_tiles], '{}/train_{:02d}_{:04d}.png'.format(flags.sample_dir, step//n_step_epoch, step))
+            tl.visualize.save_images(result.numpy(), [num_tiles, num_tiles], '{}/train_{:02d}.png'.format(flags.sample_dir, epoch))
 
 if __name__ == '__main__':
     train()