make num-image-tiles work during training as well

lucidrains · Dec 3, 2020 · 5a61846 · 5a61846
1 parent ca38892
commit 5a61846
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/stylegan2_pytorch/cli.py b/stylegan2_pytorch/cli.py
@@ -128,6 +128,7 @@ def train_from_folder(
         num_workers = num_workers,
         save_every = save_every,
         evaluate_every = evaluate_every,
+        num_image_tiles = num_image_tiles,
         trunc_psi = trunc_psi,
         fp16 = fp16,
         cl_reg = cl_reg,

diff --git a/stylegan2_pytorch/stylegan2_pytorch.py b/stylegan2_pytorch/stylegan2_pytorch.py
@@ -708,6 +708,7 @@ def __init__(
         num_workers = None,
         save_every = 1000,
         evaluate_every = 1000,
+        num_image_tiles = 8,
         trunc_psi = 0.6,
         fp16 = False,
         cl_reg = False,
@@ -764,6 +765,7 @@ def __init__(
         self.num_workers = num_workers
         self.mixed_prob = mixed_prob
 
+        self.num_image_tiles = num_image_tiles
         self.evaluate_every = evaluate_every
         self.save_every = save_every
         self.steps = 0
@@ -1054,10 +1056,10 @@ def train(self):
         self.av = None
 
     @torch.no_grad()
-    def evaluate(self, num = 0, num_image_tiles = 8, trunc = 1.0):
+    def evaluate(self, num = 0, trunc = 1.0):
         self.GAN.eval()
         ext = self.image_extension
-        num_rows = num_image_tiles
+        num_rows = self.num_image_tiles
 
         latent_dim = self.GAN.G.latent_dim
         image_size = self.GAN.G.image_size

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