offer way to turn off path length reg

lucidrains · Dec 20, 2020 · 1e4c307 · 1e4c307
1 parent 25ccfcf
commit 1e4c307
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/stylegan2_pytorch/cli.py b/stylegan2_pytorch/cli.py
@@ -95,6 +95,7 @@ def train_from_folder(
     trunc_psi = 0.75,
     mixed_prob = 0.9,
     fp16 = False,
+    no_pl_reg = False,
     cl_reg = False,
     fq_layers = [],
     fq_dict_size = 256,
@@ -131,6 +132,7 @@ def train_from_folder(
         num_image_tiles = num_image_tiles,
         trunc_psi = trunc_psi,
         fp16 = fp16,
+        no_pl_reg = no_pl_reg,
         cl_reg = cl_reg,
         fq_layers = fq_layers,
         fq_dict_size = fq_dict_size,

diff --git a/stylegan2_pytorch/stylegan2_pytorch.py b/stylegan2_pytorch/stylegan2_pytorch.py
@@ -710,6 +710,7 @@ def __init__(
         trunc_psi = 0.6,
         fp16 = False,
         cl_reg = False,
+        no_pl_reg = False,
         fq_layers = [],
         fq_dict_size = 256,
         attn_layers = [],
@@ -771,6 +772,7 @@ def __init__(
         self.av = None
         self.trunc_psi = trunc_psi
 
+        self.no_pl_reg = no_pl_reg
         self.pl_mean = None
 
         self.gradient_accumulate_every = gradient_accumulate_every
@@ -880,7 +882,7 @@ def train(self):
         aug_kwargs = {'prob': aug_prob, 'types': aug_types}
 
         apply_gradient_penalty = self.steps % 4 == 0
-        apply_path_penalty = self.steps > 5000 and self.steps % 32 == 0
+        apply_path_penalty = not self.no_pl_reg and self.steps > 5000 and self.steps % 32 == 0
         apply_cl_reg_to_generated = self.steps > 20000
 
         S = self.GAN.S if not self.is_ddp else self.S_ddp

diff --git a/stylegan2_pytorch/version.py b/stylegan2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.5.5'
+__version__ = '1.5.6'