refactor: the codes

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

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/awesome_gans/config.py b/awesome_gans/config.py
@@ -10,9 +10,11 @@ def parse_args():
     parser.add_argument('--n_channels', type=int, default=3, help='number of channel of image')
     parser.add_argument('--root_path', type=str, default='./', help='root path')
     parser.add_argument(
-        '--dataset', type=str, default='cifar10',
+        '--dataset',
+        type=str,
+        default='cifar10',
         choices=['mnist', 'cifar10', 'cifar100', 'div2k'],
-        help='type of dataset'
+        help='type of dataset',
     )
     parser.add_argument('--mnist_path', type=str, default='mnist')
     parser.add_argument('--fashion_mnist_path', type=str, default='fashion-mnist')

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, reshuffle_each_iteration=True)
+        ds = ds.shuffle(50000)
         ds = ds.batch(self.bs, drop_remainder=True)
         ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
         return ds
diff --git a/awesome_gans/wgan/config.py b/awesome_gans/wgan/config.py
@@ -9,7 +9,6 @@ def get_config():
     parser.add_argument('--epochs', default=50, type=int, help='epochs to train')
     parser.add_argument('--global_steps', default=5e4, type=int, help='iterations to train')
     parser.add_argument('--n_feats', default=64, type=int, help='number of convolution filters')
-    parser.add_argument('--lr', default=1e-4, type=float, help='learning rate for generic somethings')
     parser.add_argument('--d_lr', default=1e-4, type=float, help='learning rate of discriminator')
     parser.add_argument('--g_lr', default=1e-4, type=float, help='learning rate of generator')
     parser.add_argument(

diff --git a/awesome_gans/wgan/model.py b/awesome_gans/wgan/model.py
@@ -138,15 +138,14 @@ def train(self, dataset: tf.data.Dataset):
         for epoch in range(start_epoch, self.epochs):
             loader = tqdm(dataset, desc=f'[*] Epoch {epoch} / {self.epochs}')
             for n_iter, batch in enumerate(loader):
-                d_loss = 0.0
                 for _ in range(self.n_critics):
-                    d_loss += self.train_discriminator(batch)
+                    d_loss = self.train_discriminator(batch)
 
                 g_loss = self.train_generator()
 
                 loader.set_postfix(
-                    d_loss=f'{d_loss / self.n_critics:.6f}',
-                    g_loss=f'{g_loss:.6f}',
+                    d_loss=f'{d_loss:.5f}',
+                    g_loss=f'{g_loss:.5f}',
                 )
 
             # saving the generated samples