increase total batch size per step, and some tweaks

Author: lucidrains

bin/stylegan2_pytorch

@@ -15,11 +15,11 @@ def train_from_folder(
     image_size = 128,
     network_capacity = 16,
     transparent = False,
-    batch_size = 3,
-    gradient_accumulate_every = 5,
+    batch_size = 5,
+    gradient_accumulate_every = 6,
     num_train_steps = 150000,
     learning_rate = 2e-4,
-    ttur_mult = 2,
+    ttur_mult = 1.5,
     num_workers =  None,
     save_every = 1000,
     generate = False,

setup.py

@@ -4,7 +4,7 @@
   name = 'stylegan2_pytorch',
   packages = find_packages(),
   scripts=['bin/stylegan2_pytorch'],
-  version = '0.18.1',
+  version = '0.18.2',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

stylegan2_pytorch/stylegan2_pytorch.py

@@ -699,7 +699,7 @@ def train(self):
 
         aug_prob   = self.aug_prob
 
-        apply_gradient_penalty = self.steps > 5000 and self.steps % 4 == 0
+        apply_gradient_penalty = self.steps % 4 == 0
         apply_path_penalty = self.steps > 5000 and self.steps % 32 == 0
         apply_cl_reg_to_generated = self.steps > 20000