diff --git a/TTS/bin/train_glow_tts.py b/TTS/bin/train_glow_tts.py index 14a20149..072ad41b 100644 --- a/TTS/bin/train_glow_tts.py +++ b/TTS/bin/train_glow_tts.py @@ -497,6 +497,7 @@ def main(args): # pylint: disable=redefined-outer-name criterion = GlowTTSLoss() if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)} ...") checkpoint = torch.load(args.restore_path, map_location='cpu') try: # TODO: fix optimizer init, model.cuda() needs to be called before @@ -514,7 +515,7 @@ def main(args): # pylint: disable=redefined-outer-name for group in optimizer.param_groups: group['initial_lr'] = c.lr - print(" > Model restored from step %d" % checkpoint['step'], + print(f" > Model restored from step {checkpoint['step']:d}", flush=True) args.restore_step = checkpoint['step'] else: @@ -542,7 +543,8 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") diff --git a/TTS/bin/train_speedy_speech.py b/TTS/bin/train_speedy_speech.py index 4e521451..8e9dbc2e 100644 --- a/TTS/bin/train_speedy_speech.py +++ b/TTS/bin/train_speedy_speech.py @@ -461,6 +461,7 @@ def main(args): # pylint: disable=redefined-outer-name criterion = SpeedySpeechLoss(c) if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)} ...") checkpoint = torch.load(args.restore_path, map_location='cpu') try: # TODO: fix optimizer init, model.cuda() needs to be called before @@ -506,7 +507,8 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") diff --git a/TTS/bin/train_tacotron.py b/TTS/bin/train_tacotron.py index cdc68c94..2382c0be 100644 --- a/TTS/bin/train_tacotron.py +++ b/TTS/bin/train_tacotron.py @@ -534,12 +534,13 @@ def main(args): # pylint: disable=redefined-outer-name # setup criterion criterion = TacotronLoss(c, stopnet_pos_weight=c.stopnet_pos_weight, ga_sigma=0.4) if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)}...") checkpoint = torch.load(args.restore_path, map_location='cpu') try: - print(" > Restoring Model.") + print(" > Restoring Model...") model.load_state_dict(checkpoint['model']) # optimizer restore - print(" > Restoring Optimizer.") + print(" > Restoring Optimizer...") optimizer.load_state_dict(checkpoint['optimizer']) if "scaler" in checkpoint and c.mixed_precision: print(" > Restoring AMP Scaler...") @@ -547,7 +548,7 @@ def main(args): # pylint: disable=redefined-outer-name if c.reinit_layers: raise RuntimeError except (KeyError, RuntimeError): - print(" > Partial model initialization.") + print(" > Partial model initialization...") model_dict = model.state_dict() model_dict = set_init_dict(model_dict, checkpoint['model'], c) # torch.save(model_dict, os.path.join(OUT_PATH, 'state_dict.pt')) @@ -585,7 +586,8 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") diff --git a/TTS/bin/train_vocoder_gan.py b/TTS/bin/train_vocoder_gan.py index ecc33288..9043a560 100644 --- a/TTS/bin/train_vocoder_gan.py +++ b/TTS/bin/train_vocoder_gan.py @@ -485,6 +485,7 @@ def main(args): # pylint: disable=redefined-outer-name criterion_disc = DiscriminatorLoss(c) if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)}...") checkpoint = torch.load(args.restore_path, map_location='cpu') try: print(" > Restoring Generator Model...") @@ -523,7 +524,7 @@ def main(args): # pylint: disable=redefined-outer-name for group in optimizer_disc.param_groups: group['lr'] = c.lr_disc - print(" > Model restored from step %d" % checkpoint['step'], + print(f" > Model restored from step {checkpoint['step']:d}", flush=True) args.restore_step = checkpoint['step'] else: @@ -549,10 +550,11 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") + print(f" > Starting with best loss of {best_loss}.") keep_best = c.get('keep_best', False) keep_after = c.get('keep_after', 10000) # void if keep_best False diff --git a/TTS/bin/train_vocoder_wavegrad.py b/TTS/bin/train_vocoder_wavegrad.py index 7846aae5..271e8d4c 100644 --- a/TTS/bin/train_vocoder_wavegrad.py +++ b/TTS/bin/train_vocoder_wavegrad.py @@ -354,6 +354,7 @@ def main(args): # pylint: disable=redefined-outer-name criterion.cuda() if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)}...") checkpoint = torch.load(args.restore_path, map_location='cpu') try: print(" > Restoring Model...") @@ -397,7 +398,8 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") diff --git a/TTS/bin/train_vocoder_wavernn.py b/TTS/bin/train_vocoder_wavernn.py index 44ffef14..5fde5025 100644 --- a/TTS/bin/train_vocoder_wavernn.py +++ b/TTS/bin/train_vocoder_wavernn.py @@ -383,6 +383,7 @@ def main(args): # pylint: disable=redefined-outer-name # restore any checkpoint if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)}...") checkpoint = torch.load(args.restore_path, map_location="cpu") try: print(" > Restoring Model...") @@ -420,7 +421,8 @@ def main(args): # pylint: disable=redefined-outer-name best_loss = float('inf') print(" > Starting with inf best loss.") else: - print(args.best_path) + print(" > Restoring best loss from " + f"{os.path.basename(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}.") diff --git a/TTS/utils/arguments.py b/TTS/utils/arguments.py index d05936dc..44345dd5 100644 --- a/TTS/utils/arguments.py +++ b/TTS/utils/arguments.py @@ -169,7 +169,6 @@ def process_args(args, model_type): 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 c = load_config(args.config_path) @@ -184,8 +183,7 @@ def process_args(args, model_type): if model_class == "TTS": check_config_tts(c) elif model_class == "VOCODER": - print("Vocoder config checker not implemented, " - "skipping ...") + print("Vocoder config checker not implemented, skipping ...") else: raise ValueError(f"model type {model_type} not recognized!")