add --evaluate-every

Author: lucidrains

stylegan2_pytorch/cli.py

@@ -84,6 +84,7 @@ def train_from_folder(
     rel_disc_loss = False,
     num_workers =  None,
     save_every = 1000,
+    evaluate_every = 1000,
     generate = False,
     generate_interpolation = False,
     interpolation_num_steps = 100,
@@ -122,6 +123,7 @@ def train_from_folder(
         rel_disc_loss = rel_disc_loss,
         num_workers = num_workers,
         save_every = save_every,
+        evaluate_every = evaluate_every,
         trunc_psi = trunc_psi,
         fp16 = fp16,
         cl_reg = cl_reg,

stylegan2_pytorch/stylegan2_pytorch.py

@@ -52,7 +52,6 @@
 
 EXTS = ['jpg', 'jpeg', 'png']
 EPS = 1e-8
-EVALUATE_EVERY = 1000
 CALC_FID_NUM_IMAGES = 12800
 
 # helper classes
@@ -687,7 +686,7 @@ def forward(self, x):
         return x
 
 class Trainer():
-    def __init__(self, name, results_dir, models_dir, image_size, network_capacity, transparent = False, batch_size = 4, mixed_prob = 0.9, gradient_accumulate_every=1, lr = 2e-4, lr_mlp = 1., ttur_mult = 2, rel_disc_loss = False, num_workers = None, save_every = 1000, trunc_psi = 0.6, fp16 = False, cl_reg = False, fq_layers = [], fq_dict_size = 256, attn_layers = [], no_const = False, aug_prob = 0., aug_types = ['translation', 'cutout'], top_k_training = False, generator_top_k_gamma = 0.99, generator_top_k_frac = 0.5, dataset_aug_prob = 0., calculate_fid_every = None, is_ddp = False, rank = 0, world_size = 1, *args, **kwargs):
+    def __init__(self, name, results_dir, models_dir, image_size, network_capacity, transparent = False, batch_size = 4, mixed_prob = 0.9, gradient_accumulate_every=1, lr = 2e-4, lr_mlp = 1., ttur_mult = 2, rel_disc_loss = False, num_workers = None, save_every = 1000, evaluate_every = 1000, trunc_psi = 0.6, fp16 = False, cl_reg = False, fq_layers = [], fq_dict_size = 256, attn_layers = [], no_const = False, aug_prob = 0., aug_types = ['translation', 'cutout'], top_k_training = False, generator_top_k_gamma = 0.99, generator_top_k_frac = 0.5, dataset_aug_prob = 0., calculate_fid_every = None, is_ddp = False, rank = 0, world_size = 1, *args, **kwargs):
         self.GAN_params = [args, kwargs]
         self.GAN = None
 
@@ -719,6 +718,7 @@ def __init__(self, name, results_dir, models_dir, image_size, network_capacity,
         self.num_workers = num_workers
         self.mixed_prob = mixed_prob
 
+        self.evaluate_every = evaluate_every
         self.save_every = save_every
         self.steps = 0
 
@@ -977,8 +977,8 @@ def train(self):
             if self.steps % self.save_every == 0:
                 self.save(self.checkpoint_num)
 
-            if self.steps % EVALUATE_EVERY == 0 or (self.steps % 100 == 0 and self.steps < 2500):
-                self.evaluate(floor(self.steps / EVALUATE_EVERY))
+            if self.steps % self.evaluate_every == 0 or (self.steps % 100 == 0 and self.steps < 2500):
+                self.evaluate(floor(self.steps / self.evaluate_every))
 
             if exists(self.calculate_fid_every) and self.steps % self.calculate_fid_every == 0 and self.steps != 0:
                 num_batches = math.ceil(CALC_FID_NUM_IMAGES / self.batch_size)

stylegan2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.2.5'
+__version__ = '1.2.6'