Remove thrash and allow to load latest pkl files

andersonfaaria · Aug 7, 2020 · 756bae3 · 756bae3
1 parent 212ba54
commit 756bae3
Show file tree

Hide file tree

Showing 9 changed files with 4 additions and 91 deletions.
diff --git a/gen.bat b/gen.bat
diff --git a/play_dlatents.bat b/play_dlatents.bat
diff --git a/play_vectors.bat b/play_vectors.bat
diff --git a/prepare_dataset.bat b/prepare_dataset.bat
diff --git a/project.bat b/project.bat
diff --git a/reduce_model.bat b/reduce_model.bat
diff --git a/src/training/training_loop.py b/src/training/training_loop.py
@@ -130,7 +130,7 @@ def training_loop(
     network_snapshot_ticks  = 50,       # How often to save network snapshots? None = only save 'networks-final.pkl'.
     save_tf_graph           = False,    # Include full TensorFlow computation graph in the tfevents file?
     save_weight_histograms  = False,    # Include weight histograms in the tfevents file?
-    resume_pkl              = None,     # Network pickle to resume training from, None = train from scratch.
+    resume_pkl              = 'latest',     # Network pickle to resume training from, None = train from scratch.
     resume_kimg             = 0.0,      # Assumed training progress at the beginning. Affects reporting and training schedule.
     resume_time             = 0.0,      # Assumed wallclock time at the beginning. Affects reporting.
     restore_partial_fn      = None,   # Filename of network for partial restore
@@ -162,7 +162,9 @@ def training_loop(
             D = tflib.Network('D', num_channels=training_set.shape[0], resolution=resolution, label_size=training_set.label_size, **D_args)
             Gs = G.clone('Gs')
         if resume_pkl is not None:
-            if resume_pkl == 'restore_partial':
+            if resume_pkl == 'latest':
+                resume_pkl, resume_kimg = misc.locate_latest_pkl(dnnlib.submit_config.run_dir_root)
+            elif resume_pkl == 'restore_partial':
                 print(' Restore partially...')
                 # Initialize networks
                 G = tflib.Network('G', num_channels=training_set.shape[0], resolution=resolution, label_size=training_set.label_size, **G_args)

diff --git a/train.bat b/train.bat
diff --git a/train_resume.bat b/train_resume.bat