Update README and better control when to save image snapshot of fakes…

…, as well as its resolution; set up more configurations for GPUs
PDillis · Jul 29, 2021 · b9db34f · b9db34f
1 parent 8a05574
commit b9db34f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -5,6 +5,8 @@ This README will eventually be updated to correctly show how to run the followin
   a need to specify where the output will be saved. The output directories will be automatically generated, as was done in the 
   previous implementations ([ProGAN](https://github.com/tkarras/progressive_growing_of_gans) / [StyleGAN](https://github.com/NVlabs/stylegan) 
   / [StyleGAN2](https://github.com/NVlabs/stylegan2)).
+* **Setup of image snapshot resolution and its saving frequency**: Use `--img-snap` to set how often to save the snapshot grid of
+  `fakes.jpg`, as well as its resolution via `--snap-res` (options: `1080p`, `4k`, and `8k`).
 * **Random interpolation video**: Run `python generate.py random-video --help` on how to run it, but usage is similar as in 
   [StyleGAN2](https://github.com/PDillis/stylegan2-fun#random-interpolation) (only the name changes).
 * **Style mixing videos**: Run `python style_mixing.py video --help` on how to run it, but usage is similar as in 

diff --git a/train.py b/train.py
@@ -33,6 +33,7 @@ def setup_training_loop_kwargs(
     # General options (not included in desc).
     gpus       = None, # Number of GPUs: <int>, default = 1 gpu
     snap       = None, # Snapshot interval: <int>, default = 50 ticks
+    img_snap   = None, # Image snapshot interval: <int>, default = 50 ticks
     snap_res   = None,  # Snapshot size (1080p/4k/8k)
     metrics    = None, # List of metric names: [], ['fid50k_full'] (default), ...
     seed       = None, # Random seed: <int>, default = 0
@@ -86,9 +87,13 @@ def setup_training_loop_kwargs(
     assert isinstance(snap, int)
     if snap < 1:
         raise UserError('--snap must be at least 1')
-    args.image_snapshot_ticks = 1  # I prefer to see how it evolves vs. find out only when the model saves
     args.network_snapshot_ticks = snap
 
+    assert isinstance(img_snap, int)
+    if img_snap < 1:
+        raise UserError('--img-snap must be at least 1')
+    args.image_snapshot_ticks = img_snap  # I prefer to see how it evolves vs. find out only when the model saves
+
     args.snap_res = '8k' if snap_res is None else snap_res
 
     if metrics is None:
@@ -171,6 +176,8 @@ def setup_training_loop_kwargs(
         '24gb-4gpu':      dict(ref_gpus=4,  kimg=25000,  mb=48, mbstd=12, fmaps=1,   lrate=0.002,  gamma=10,   ema=10,  ramp=None, map=8),  # Made for 1024x1024 dataset; adapted from @dvschultz
         '24gb-2gpu-cplx': dict(ref_gpus=2,  kimg=25000,  mb=24, mbstd=12, fmaps=1,   lrate=0.002,  gamma=10,   ema=10,  ramp=None, map=4),  # Complex model, inspired by @aydao
         '24gb-4gpu-cplx': dict(ref_gpus=4,  kimg=25000,  mb=48, mbstd=12, fmaps=1,   lrate=0.002,  gamma=10,   ema=10,  ramp=None, map=4),  # Complex model, inspired by @aydao
+        '48gb-2gpu':      dict(ref_gpus=2,  kimg=25000,  mb=48, mbstd=24, fmaps=1,   lrate=0.002,  gamma=10,   ema=10,  ramp=None, map=8),  # Made for 1024x1024 dataset; adapted from @dvschultz
+        '48gb-4gpu':      dict(ref_gpus=4,  kimg=25000,  mb=96, mbstd=24, fmaps=1,   lrate=0.002,  gamma=10,   ema=10,  ramp=None, map=8),  # Made for 1024x1024 dataset; adapted from @dvschultz
     }
 
     assert cfg in cfg_specs
@@ -291,7 +298,6 @@ def setup_training_loop_kwargs(
             raise UserError('--augpipe cannot be specified with --aug=noaug')
         desc += f'-{augpipe}'
 
-    # Remove yflip=1; move to other configs using yflip (blit-y, bg-y, bgc-y, ...)
     augpipe_specs = {
         'blit':   dict(xflip=1, rotate90=1, xint=1),
         'geom':   dict(scale=1, rotate=1, aniso=1, xfrac=1),
@@ -444,6 +450,7 @@ def convert(self, value, param, ctx):
 @click.option('--outdir', help='Where to save the results', type=click.Path(file_okay=False), default=os.path.join(os.getcwd(), 'training-runs'), show_default=True, metavar='DIR')
 @click.option('--gpus', help='Number of GPUs to use', type=click.IntRange(min=1, max=8), default=1, show_default=True, metavar='INT')
 @click.option('--snap', help='Snapshot interval in ticks', type=int, default=50, show_default=True, metavar='INT')
+@click.option('--img-snap', help='Image snapshot interval', type=int, default=50, show_default=True, metavar='INT')
 @click.option('--snap-res', help='Image snapshot resolution', type=click.Choice(['1080p', '4k', '8k']), default='8k', show_default=True)
 @click.option('--metrics', help='Comma-separated list of metrics or "none" [default: fid50k_full]', type=CommaSeparatedList(), )
 @click.option('--seed', help='Random seed [default: 0]', type=int, metavar='INT')