Medix-AI/coqui-tts
Medix-AI
/
coqui-tts
mirror of https://github.com/coqui-ai/TTS.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

brushed up printing model load path and best loss path

This commit is contained in:
gerazov 2021-02-12 10:55:45 +01:00 committed by Eren Gölge
parent f2e474cd37
commit 2db40457e8
7 changed files with 25 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
TTS/bin/train_glow_tts.py
View File

@ -500,6 +500,7 @@ def main(args): # pylint: disable=redefined-outer-name
criterion = GlowTTSLoss() criterion = GlowTTSLoss()
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)} ...")
checkpoint = torch.load(args.restore_path, map_location='cpu') checkpoint = torch.load(args.restore_path, map_location='cpu')
try: try:
# TODO: fix optimizer init, model.cuda() needs to be called before # TODO: fix optimizer init, model.cuda() needs to be called before
@ -517,7 +518,7 @@ def main(args): # pylint: disable=redefined-outer-name
for group in optimizer.param_groups: for group in optimizer.param_groups:
group['initial_lr'] = c.lr group['initial_lr'] = c.lr
print(" > Model restored from step %d" % checkpoint['step'], print(f" > Model restored from step {checkpoint['step']:d}",
flush=True) flush=True)
args.restore_step = checkpoint['step'] args.restore_step = checkpoint['step']
else: else:
@ -545,7 +546,8 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] map_location='cpu')['model_loss']
print(f" > Starting with loaded last best loss {best_loss}.") print(f" > Starting with loaded last best loss {best_loss}.")

4
TTS/bin/train_speedy_speech.py
View File

@ -464,6 +464,7 @@ def main(args): # pylint: disable=redefined-outer-name
criterion = SpeedySpeechLoss(c) criterion = SpeedySpeechLoss(c)
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)} ...")
checkpoint = torch.load(args.restore_path, map_location='cpu') checkpoint = torch.load(args.restore_path, map_location='cpu')
try: try:
# TODO: fix optimizer init, model.cuda() needs to be called before # TODO: fix optimizer init, model.cuda() needs to be called before
@ -509,7 +510,8 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] map_location='cpu')['model_loss']
print(f" > Starting with loaded last best loss {best_loss}.") print(f" > Starting with loaded last best loss {best_loss}.")

10
TTS/bin/train_tacotron.py
View File

@ -538,12 +538,13 @@ def main(args): # pylint: disable=redefined-outer-name
# setup criterion # setup criterion
criterion = TacotronLoss(c, stopnet_pos_weight=c.stopnet_pos_weight, ga_sigma=0.4) criterion = TacotronLoss(c, stopnet_pos_weight=c.stopnet_pos_weight, ga_sigma=0.4)
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)}...")
checkpoint = torch.load(args.restore_path, map_location='cpu') checkpoint = torch.load(args.restore_path, map_location='cpu')
try: try:
print(" > Restoring Model.") print(" > Restoring Model...")
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
# optimizer restore # optimizer restore
print(" > Restoring Optimizer.") print(" > Restoring Optimizer...")
optimizer.load_state_dict(checkpoint['optimizer']) optimizer.load_state_dict(checkpoint['optimizer'])
if "scaler" in checkpoint and c.mixed_precision: if "scaler" in checkpoint and c.mixed_precision:
print(" > Restoring AMP Scaler...") print(" > Restoring AMP Scaler...")
@ -551,7 +552,7 @@ def main(args): # pylint: disable=redefined-outer-name
if c.reinit_layers: if c.reinit_layers:
raise RuntimeError raise RuntimeError
except (KeyError, RuntimeError): except (KeyError, RuntimeError):
print(" > Partial model initialization.") print(" > Partial model initialization...")
model_dict = model.state_dict() model_dict = model.state_dict()
model_dict = set_init_dict(model_dict, checkpoint['model'], c) model_dict = set_init_dict(model_dict, checkpoint['model'], c)
# torch.save(model_dict, os.path.join(OUT_PATH, 'state_dict.pt')) # torch.save(model_dict, os.path.join(OUT_PATH, 'state_dict.pt'))
@ -589,7 +590,8 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] map_location='cpu')['model_loss']
print(f" > Starting with loaded last best loss {best_loss}.") print(f" > Starting with loaded last best loss {best_loss}.")

8
TTS/bin/train_vocoder_gan.py
View File

@ -485,6 +485,7 @@ def main(args): # pylint: disable=redefined-outer-name
criterion_disc = DiscriminatorLoss(c) criterion_disc = DiscriminatorLoss(c)
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)}...")
checkpoint = torch.load(args.restore_path, map_location='cpu') checkpoint = torch.load(args.restore_path, map_location='cpu')
try: try:
print(" > Restoring Generator Model...") print(" > Restoring Generator Model...")
@ -523,7 +524,7 @@ def main(args): # pylint: disable=redefined-outer-name
for group in optimizer_disc.param_groups: for group in optimizer_disc.param_groups:
group['lr'] = c.lr_disc group['lr'] = c.lr_disc
print(" > Model restored from step %d" % checkpoint['step'], print(f" > Model restored from step {checkpoint['step']:d}",
flush=True) flush=True)
args.restore_step = checkpoint['step'] args.restore_step = checkpoint['step']
else: else:
@ -549,10 +550,11 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] 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_best = c.get('keep_best', False)
keep_after = c.get('keep_after', 10000) # void if keep_best False keep_after = c.get('keep_after', 10000) # void if keep_best False

4
TTS/bin/train_vocoder_wavegrad.py
View File

@ -354,6 +354,7 @@ def main(args): # pylint: disable=redefined-outer-name
criterion.cuda() criterion.cuda()
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)}...")
checkpoint = torch.load(args.restore_path, map_location='cpu') checkpoint = torch.load(args.restore_path, map_location='cpu')
try: try:
print(" > Restoring Model...") print(" > Restoring Model...")
@ -397,7 +398,8 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] map_location='cpu')['model_loss']
print(f" > Starting with loaded last best loss {best_loss}.") print(f" > Starting with loaded last best loss {best_loss}.")

4
TTS/bin/train_vocoder_wavernn.py
View File

@ -383,6 +383,7 @@ def main(args): # pylint: disable=redefined-outer-name
# restore any checkpoint # restore any checkpoint
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)}...")
checkpoint = torch.load(args.restore_path, map_location="cpu") checkpoint = torch.load(args.restore_path, map_location="cpu")
try: try:
print(" > Restoring Model...") print(" > Restoring Model...")
@ -420,7 +421,8 @@ def main(args): # pylint: disable=redefined-outer-name
best_loss = float('inf') best_loss = float('inf')
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(args.best_path) print(" > Restoring best loss from "
f"{os.path.basename(args.best_path)} ...")
best_loss = torch.load(args.best_path, best_loss = torch.load(args.best_path,
map_location='cpu')['model_loss'] map_location='cpu')['model_loss']
print(f" > Starting with loaded last best loss {best_loss}.") print(f" > Starting with loaded last best loss {best_loss}.")

4
TTS/utils/arguments.py
View File

@ -157,7 +157,6 @@ def process_args(args, model_type):
args.restore_path, best_model = get_last_models(args.continue_path) args.restore_path, best_model = get_last_models(args.continue_path)
if not args.best_path: if not args.best_path:
args.best_path = best_model args.best_path = best_model
print(f" > Training continues for {args.restore_path}")
# setup output paths and read configs # setup output paths and read configs
c = load_config(args.config_path) c = load_config(args.config_path)
@ -171,8 +170,7 @@ def process_args(args, model_type):
if model_class == "TTS": if model_class == "TTS":
check_config_tts(c) check_config_tts(c)
elif model_class == "VOCODER": elif model_class == "VOCODER":
print("Vocoder config checker not implemented, " print("Vocoder config checker not implemented, skipping ...")
"skipping ...")
else: else:
raise ValueError(f"model type {model_type} not recognized!") raise ValueError(f"model type {model_type} not recognized!")