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

restore schedulers only if training is continuing a previous training 

inherit nn.Module for TorchSTFT
This commit is contained in:
Eren Gölge 2021-04-09 19:31:28 +02:00
parent 2c71c6d8cd
commit 5b70da2e3f
2 changed files with 15 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
TTS/bin/train_vocoder_gan.py
View File

@ -515,17 +515,19 @@ def main(args): # pylint: disable=redefined-outer-name
model_disc.load_state_dict(checkpoint['model_disc'])
print(" > Restoring Discriminator Optimizer...")
optimizer_disc.load_state_dict(checkpoint['optimizer_disc'])
if 'scheduler' in checkpoint and scheduler_gen is not None:
print(" > Restoring Generator LR Scheduler...")
scheduler_gen.load_state_dict(checkpoint['scheduler'])
# NOTE: Not sure if necessary
scheduler_gen.optimizer = optimizer_gen
if 'scheduler_disc' in checkpoint and scheduler_disc is not None:
print(" > Restoring Discriminator LR Scheduler...")
scheduler_disc.load_state_dict(checkpoint['scheduler_disc'])
scheduler_disc.optimizer = optimizer_disc
if c.lr_scheduler_disc == "ExponentialLR":
scheduler_disc.last_epoch = checkpoint['epoch']
# restore schedulers if it is a continuing training.
if args.continue_path != '':
if 'scheduler' in checkpoint and scheduler_gen is not None:
print(" > Restoring Generator LR Scheduler...")
scheduler_gen.load_state_dict(checkpoint['scheduler'])
# NOTE: Not sure if necessary
scheduler_gen.optimizer = optimizer_gen
if 'scheduler_disc' in checkpoint and scheduler_disc is not None:
print(" > Restoring Discriminator LR Scheduler...")
scheduler_disc.load_state_dict(checkpoint['scheduler_disc'])
scheduler_disc.optimizer = optimizer_disc
if c.lr_scheduler_disc == "ExponentialLR":
scheduler_disc.last_epoch = checkpoint['epoch']
except RuntimeError:
# restore only matching layers.
print(" > Partial model initialization...")

3
TTS/vocoder/layers/losses.py
View File

@ -4,7 +4,7 @@ from torch import nn
from torch.nn import functional as F
class TorchSTFT():
class TorchSTFT(nn.Module): # pylint: disable=abstract-method
"""TODO: Merge this with audio.py"""
def __init__(self,
n_fft,
@ -34,6 +34,7 @@ class TorchSTFT():
if use_mel:
self._build_mel_basis()
@torch.no_grad()
def __call__(self, x):
"""Compute spectrogram frames by torch based stft.