loading last checkpoint/best_model works, deleting last best models options added, loading last best_loss added

TTS/bin/train_glow_tts.py

@@ -538,8 +538,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_params = count_parameters(model)
     print("\n > Model has {} parameters".format(num_params), flush=True)
 
-    if 'best_loss' not in locals():
+    if args.restore_step == 0 or not args.best_path:
         best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     # define dataloaders
     train_loader = setup_loader(ap, 1, is_val=False, verbose=True)
@@ -549,7 +557,8 @@ def main(args):  # pylint: disable=redefined-outer-name
     model = data_depended_init(train_loader, model)
     for epoch in range(0, c.epochs):
         c_logger.print_epoch_start(epoch, c.epochs)
-        train_avg_loss_dict, global_step = train(train_loader, model, criterion, optimizer,
+        train_avg_loss_dict, global_step = train(train_loader, model,
+                                                 criterion, optimizer,
                                                  scheduler, ap, global_step,
                                                  epoch)
         eval_avg_loss_dict = evaluate(eval_loader, model, criterion, ap,
@@ -558,8 +567,9 @@ def main(args):  # pylint: disable=redefined-outer-name
         target_loss = train_avg_loss_dict['avg_loss']
         if c.run_eval:
             target_loss = eval_avg_loss_dict['avg_loss']
-        best_loss = save_best_model(target_loss, best_loss, model, optimizer, global_step, epoch, c.r,
-                                    OUT_PATH)
+        best_loss = save_best_model(target_loss, best_loss, model, optimizer,
+                                    global_step, epoch, c.r, OUT_PATH,
+                                    keep_best=keep_best, keep_after=keep_after)
 
 
 if __name__ == '__main__':

TTS/bin/train_speedy_speech.py

@@ -502,8 +502,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_params = count_parameters(model)
     print("\n > Model has {} parameters".format(num_params), flush=True)
 
-    if 'best_loss' not in locals():
+    if args.restore_step == 0 or not args.best_path:
         best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     # define dataloaders
     train_loader = setup_loader(ap, 1, is_val=False, verbose=True)
@@ -522,8 +530,8 @@ def main(args):  # pylint: disable=redefined-outer-name
         if c.run_eval:
             target_loss = eval_avg_loss_dict['avg_loss']
         best_loss = save_best_model(target_loss, best_loss, model, optimizer,
-                                    global_step, epoch, c.r,
-                                    OUT_PATH)
+                                    global_step, epoch, c.r, OUT_PATH,
+                                    keep_best=keep_best, keep_after=keep_after)
 
 
 if __name__ == '__main__':

TTS/bin/train_tacotron.py

@@ -581,8 +581,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_params = count_parameters(model)
     print("\n > Model has {} parameters".format(num_params), flush=True)
 
-    if 'best_loss' not in locals():
+    if args.restore_step == 0 or not args.best_path:
         best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     # define data loaders
     train_loader = setup_loader(ap,
@@ -634,6 +642,8 @@ def main(args):  # pylint: disable=redefined-outer-name
             epoch,
             c.r,
             OUT_PATH,
+            keep_best=keep_best,
+            keep_after=keep_after,
             scaler=scaler.state_dict() if c.mixed_precision else None
         )
 

TTS/bin/train_vocoder_gan.py

@@ -545,8 +545,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_params = count_parameters(model_disc)
     print(" > Discriminator has {} parameters".format(num_params), flush=True)
 
-    if 'best_loss' not in locals():
+    if args.restore_step == 0 or not args.best_path:
         best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     global_step = args.restore_step
     for epoch in range(0, c.epochs):
@@ -571,7 +579,10 @@ def main(args):  # pylint: disable=redefined-outer-name
                                     global_step,
                                     epoch,
                                     OUT_PATH,
-                                    model_losses=eval_avg_loss_dict)
+                                    keep_best=keep_best,
+                                    keep_after=keep_after,
+                                    model_losses=eval_avg_loss_dict,
+                                    )
 
 
 if __name__ == '__main__':

TTS/bin/train_vocoder_wavegrad.py

@@ -393,8 +393,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_params = count_parameters(model)
     print(" > WaveGrad has {} parameters".format(num_params), flush=True)
 
-    if 'best_loss' not in locals():
+    if args.restore_step == 0 or not args.best_path:
         best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     global_step = args.restore_step
     for epoch in range(0, c.epochs):
@@ -416,6 +424,8 @@ def main(args):  # pylint: disable=redefined-outer-name
             global_step,
             epoch,
             OUT_PATH,
+            keep_best=keep_best,
+            keep_after=keep_after,
             model_losses=eval_avg_loss_dict,
             scaler=scaler.state_dict() if c.mixed_precision else None
         )

TTS/bin/train_vocoder_wavernn.py

@@ -416,8 +416,16 @@ def main(args):  # pylint: disable=redefined-outer-name
     num_parameters = count_parameters(model_wavernn)
     print(" > Model has {} parameters".format(num_parameters), flush=True)
 
-    if "best_loss" not in locals():
-        best_loss = float("inf")
+    if args.restore_step == 0 or not args.best_path:
+        best_loss = float('inf')
+        print(" > Starting with inf best loss.")
+    else:
+        print(args.best_path)
+        best_loss = torch.load(args.best_path,
+                               map_location='cpu')['model_loss']
+        print(f" > Starting with loaded last best loss {best_loss}.")
+    keep_best = c.get('keep_best', False)
+    keep_after = c.get('keep_after', 10000)  # void if keep_best False
 
     global_step = args.restore_step
     for epoch in range(0, c.epochs):
@@ -440,6 +448,8 @@ def main(args):  # pylint: disable=redefined-outer-name
             global_step,
             epoch,
             OUT_PATH,
+            keep_best=keep_best,
+            keep_after=keep_after,
             model_losses=eval_avg_loss_dict,
             scaler=scaler.state_dict() if c.mixed_precision else None
         )

TTS/tts/configs/config.json

@@ -121,6 +121,8 @@
     "print_eval": false,     // If True, it prints intermediate loss values in evalulation.
     "save_step": 10000,      // Number of training steps expected to save traninpg stats and checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/tts/configs/glow_tts_gated_conv.json

@@ -93,6 +93,8 @@
     "print_eval": false,     // If True, it prints intermediate loss values in evalulation.
     "save_step": 5000,      // Number of training steps expected to save traninpg stats and checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
     "apex_amp_level": null,
 

TTS/tts/configs/glow_tts_ljspeech.json

@@ -105,6 +105,8 @@
         "print_eval": false,     // If True, it prints intermediate loss values in evalulation.
         "save_step": 5000,      // Number of training steps expected to save traninpg stats and checkpoints.
         "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+        "keep_best": false,  // If true, keeps all best_models after keep_after steps
+        "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
         "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
         // DATA LOADING

TTS/tts/configs/ljspeech_tacotron2_dynamic_conv_attn.json

@@ -121,6 +121,8 @@
     "print_eval": false,     // If True, it prints intermediate loss values in evalulation.
     "save_step": 10000,      // Number of training steps expected to save traninpg stats and checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/tts/configs/speedy_speech_ljspeech.json

@@ -109,6 +109,8 @@
    "print_eval": false,     // If True, it prints intermediate loss values in evalulation.
    "save_step": 5000,      // Number of training steps expected to save traninpg stats and checkpoints.
    "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+   "keep_best": false,  // If true, keeps all best_models after keep_after steps
+   "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
    "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.:set n
    "mixed_precision": false,
 

TTS/utils/arguments.py

@@ -43,6 +43,11 @@ def parse_arguments(argv):
         type=str,
         help="Model file to be restored. Use to finetune a model.",
         default="")
+    parser.add_argument(
+        "--best_path",
+        type=str,
+        help="Best model file to be used for extracting best loss.",
+        default="")
     parser.add_argument(
         "--config_path",
         type=str,
@@ -67,11 +72,11 @@ def parse_arguments(argv):
     return parser.parse_args()
 
 
-def get_last_checkpoint(path):
-    """Get latest checkpoint from a list of filenames.
+def get_last_models(path):
+    """Get latest checkpoint or/and best model in path.
 
     It is based on globbing for `*.pth.tar` and the RegEx
-    `checkpoint_([0-9]+)`.
+    `(checkpoint|best_model)_([0-9]+)`.
 
     Parameters
     ----------
@@ -81,7 +86,7 @@ def get_last_checkpoint(path):
     Raises
     ------
     ValueError
-        If no checkpoint files are found.
+        If no checkpoint or best_model files are found.
 
     Returns
     -------
@@ -89,22 +94,37 @@ def get_last_checkpoint(path):
         Last checkpoint filename.
 
     """
-    last_checkpoint_num = 0
-    last_checkpoint = None
-    filenames = glob.glob(
-        os.path.join(path, "/*.pth.tar"))
-    for filename in filenames:
+    file_names = glob.glob(os.path.join(path, "*.pth.tar"))
+    last_models = {}
+    last_model_nums = {}
+    for key in ['checkpoint', 'best_model']:
+        last_model_num = 0
+        last_model = None
+        for file_name in file_names:
             try:
-            checkpoint_num = int(
-                re.search(r"checkpoint_([0-9]+)", filename).groups()[0])
-            if checkpoint_num > last_checkpoint_num:
-                last_checkpoint_num = checkpoint_num
-                last_checkpoint = filename
+                model_num = int(re.search(
+                    f"{key}_([0-9]+)", file_name).groups()[0])
+                if model_num > last_model_num:
+                    last_model_num = model_num
+                    last_model = file_name
             except AttributeError:  # if there's no match in the filename
-            pass
-    if last_checkpoint is None:
-        raise ValueError(f"No checkpoints in {path}!")
-    return last_checkpoint
+                continue
+        last_models[key] = last_model
+        last_model_nums[key] = last_model_num
+
+    # check what models were found
+    if not last_models:
+        raise ValueError(f"No models found in continue path {path}!")
+    elif 'checkpoint' not in last_models:  # no checkpoint just best model
+        last_models['checkpoint'] = last_models['best_model']
+    elif 'best_model' not in last_models:  # no best model
+        # this shouldn't happen, but let's handle it just in case
+        last_models['best_model'] = None
+    # finally check if last best model is more recent than checkpoint
+    elif last_model_nums['best_model'] > last_model_nums['checkpoint']:
+        last_models['checkpoint'] = last_models['best_model']
+
+    return last_models['checkpoint'], last_models['best_model']
 
 
 def process_args(args, model_type):
@@ -143,15 +163,12 @@ def process_args(args, model_type):
         Class that does the TensorBoard loggind.
 
     """
-    if args.continue_path != "":
+    if args.continue_path:
         args.output_path = args.continue_path
         args.config_path = os.path.join(args.continue_path, "config.json")
-        list_of_files = glob.glob(
-            os.path.join(args.continue_path, "*.pth.tar")
-            )  # * means all if need specific format then *.csv
-        args.restore_path = max(list_of_files, key=os.path.getctime)
-        # checkpoint number based continuing
-        # args.restore_path = get_last_checkpoint(args.continue_path)
+        args.restore_path, best_model = get_last_models(args.continue_path)
+        if not args.best_path:
+            args.best_path = best_model
         print(f" > Training continues for {args.restore_path}")
 
     # setup output paths and read configs
@@ -178,7 +195,7 @@ def process_args(args, model_type):
         print("   >  Mixed precision mode is ON")
 
     out_path = args.continue_path
-    if args.continue_path == "":
+    if not out_path:
         out_path = create_experiment_folder(c.output_path, c.run_name,
                                             args.debug)
 

TTS/vocoder/configs/multiband-melgan_and_rwd_config.json

@@ -138,6 +138,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 25000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/multiband_melgan_config.json

@@ -128,6 +128,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 25000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/multiband_melgan_config_mozilla.json

@@ -141,6 +141,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 25000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/parallel_wavegan_config.json

@@ -130,6 +130,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 25000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/universal_fullband_melgan.json

@@ -124,6 +124,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 25000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/wavegrad_libritts.json

@@ -103,6 +103,8 @@
     "print_eval": false,     // If True, it prints loss values for each step in eval run.
     "save_step": 5000,      // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true,     // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": true,     // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
     // DATA LOADING

TTS/vocoder/configs/wavernn_config.json

@@ -89,6 +89,8 @@
     "print_eval": false, // If True, it prints loss values for each step in eval run.
     "save_step": 25000, // Number of training steps expected to plot training stats on TB and save model checkpoints.
     "checkpoint": true, // If true, it saves checkpoints per "save_step"
+    "keep_best": false,  // If true, keeps all best_models after keep_after steps
+    "keep_after": 10000,    // Global step after which to keep best models if keep_best is true
     "tb_model_param_stats": false, // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
 
 // DATA LOADING

TTS/vocoder/utils/io.py

@@ -1,4 +1,5 @@
 import os
+import glob
 import torch
 import datetime
 import pickle as pickle_tts
@@ -61,12 +62,13 @@ def save_checkpoint(model, optimizer, scheduler, model_disc, optimizer_disc,
                scheduler_disc, current_step, epoch, checkpoint_path, **kwargs)
 
 
-def save_best_model(target_loss, best_loss, model, optimizer, scheduler,
+def save_best_model(current_loss, best_loss, model, optimizer, scheduler,
                     model_disc, optimizer_disc, scheduler_disc, current_step,
-                    epoch, output_folder, **kwargs):
-    if target_loss < best_loss:
-        file_name = 'best_model.pth.tar'
-        checkpoint_path = os.path.join(output_folder, file_name)
+                    epoch, out_path, keep_best=False, keep_after=10000,
+                    **kwargs):
+    if current_loss < best_loss:
+        best_model_name = f'best_model_{current_step}.pth.tar'
+        checkpoint_path = os.path.join(out_path, best_model_name)
         print(" > BEST MODEL : {}".format(checkpoint_path))
         save_model(model,
                    optimizer,
@@ -77,7 +79,21 @@ def save_best_model(target_loss, best_loss, model, optimizer, scheduler,
                    current_step,
                    epoch,
                    checkpoint_path,
-                   model_loss=target_loss,
+                   model_loss=current_loss,
                    **kwargs)
-        best_loss = target_loss
+        # only delete previous if current is saved successfully
+        if not keep_best or (current_step < keep_after):
+            model_names = glob.glob(
+                os.path.join(out_path, 'best_model*.pth.tar'))
+            for model_name in model_names:
+                if os.path.basename(model_name) == best_model_name:
+                    continue
+                os.remove(model_name)
+        # create symlink to best model for convinience
+        link_name = 'best_model.pth.tar'
+        link_path = os.path.join(out_path, link_name)
+        if os.path.islink(link_path) or os.path.isfile(link_path):
+            os.remove(link_path)
+        os.symlink(best_model_name, os.path.join(out_path, link_name))
+        best_loss = current_loss
     return best_loss