updated to current dev

This commit is contained in:
gerazov 2021-02-06 22:59:52 +01:00 committed by Eren Gölge
parent 6f06e31541
commit 10d5a63d49
6 changed files with 109 additions and 107 deletions

View File

@ -116,7 +116,7 @@ def format_data(data):
avg_text_length, avg_spec_length, attn_mask, item_idx
def data_depended_init(data_loader, model, ap):
def data_depended_init(data_loader, model):
"""Data depended initialization for activation normalization."""
if hasattr(model, 'module'):
for f in model.module.decoder.flows:
@ -135,7 +135,7 @@ def data_depended_init(data_loader, model, ap):
# format data
text_input, text_lengths, mel_input, mel_lengths, spekaer_embed,\
_, _, attn_mask, item_idx = format_data(data)
_, _, attn_mask, _ = format_data(data)
# forward pass model
_ = model.forward(
@ -174,7 +174,7 @@ def train(data_loader, model, criterion, optimizer, scheduler,
# format data
text_input, text_lengths, mel_input, mel_lengths, speaker_c,\
avg_text_length, avg_spec_length, attn_mask, item_idx = format_data(data)
avg_text_length, avg_spec_length, attn_mask, _ = format_data(data)
loader_time = time.time() - end_time
@ -329,7 +329,7 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
# format data
text_input, text_lengths, mel_input, mel_lengths, speaker_c,\
_, _, attn_mask, item_idx = format_data(data)
_, _, attn_mask, _ = format_data(data)
# forward pass model
z, logdet, y_mean, y_log_scale, alignments, o_dur_log, o_total_dur = model.forward(
@ -546,13 +546,14 @@ def main(args): # pylint: disable=redefined-outer-name
eval_loader = setup_loader(ap, 1, is_val=True, verbose=True)
global_step = args.restore_step
model = data_depended_init(train_loader, model, ap)
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,
scheduler, ap, global_step,
epoch)
eval_avg_loss_dict = evaluate(eval_loader , model, criterion, ap, global_step, epoch)
eval_avg_loss_dict = evaluate(eval_loader , model, criterion, ap,
global_step, epoch)
c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
target_loss = train_avg_loss_dict['avg_loss']
if c.run_eval:

View File

@ -515,12 +515,14 @@ def main(args): # pylint: disable=redefined-outer-name
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, global_step, epoch)
eval_avg_loss_dict = evaluate(eval_loader , model, criterion, ap,
global_step, epoch)
c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
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,
best_loss = save_best_model(target_loss, best_loss, model, optimizer,
global_step, epoch, c.r,
OUT_PATH)

View File

@ -9,8 +9,8 @@ from random import randrange
import numpy as np
import torch
from TTS.utils.arguments import parse_arguments, process_args
from torch.utils.data import DataLoader
from TTS.utils.arguments import parse_arguments, process_args
from TTS.tts.datasets.preprocess import load_meta_data
from TTS.tts.datasets.TTSDataset import MyDataset
from TTS.tts.layers.losses import TacotronLoss
@ -181,8 +181,8 @@ def train(data_loader, model, criterion, optimizer, optimizer_st, scheduler,
loss_dict = criterion(postnet_output, decoder_output, mel_input,
linear_input, stop_tokens, stop_targets,
mel_lengths, decoder_backward_output,
alignments, alignment_lengths, alignments_backward,
text_lengths)
alignments, alignment_lengths,
alignments_backward, text_lengths)
# check nan loss
if torch.isnan(loss_dict['loss']).any():
@ -534,8 +534,7 @@ def main(args): # pylint: disable=redefined-outer-name
optimizer_st = None
# setup criterion
criterion = TacotronLoss(c, stopnet_pos_weight=10.0, ga_sigma=0.4)
criterion = TacotronLoss(c, stopnet_pos_weight=c.stopnet_pos_weight, ga_sigma=0.4)
if args.restore_path:
checkpoint = torch.load(args.restore_path, map_location='cpu')
try:
@ -637,7 +636,8 @@ def main(args): # pylint: disable=redefined-outer-name
epoch,
c.r,
OUT_PATH,
scaler=scaler.state_dict() if c.mixed_precision else None)
scaler=scaler.state_dict() if c.mixed_precision else None
)
if __name__ == '__main__':

View File

@ -8,8 +8,8 @@ import traceback
from inspect import signature
import torch
from TTS.utils.arguments import parse_arguments, process_args
from torch.utils.data import DataLoader
from TTS.utils.arguments import parse_arguments, process_args
from TTS.utils.audio import AudioProcessor
from TTS.utils.generic_utils import (KeepAverage, count_parameters,
remove_experiment_folder, set_init_dict)
@ -33,9 +33,8 @@ use_cuda, num_gpus = setup_torch_training_env(True, True)
def setup_loader(ap, is_val=False, verbose=False):
if is_val and not c.run_eval:
loader = None
else:
if not is_val or c.run_eval:
dataset = GANDataset(ap=ap,
items=eval_data if is_val else train_data,
seq_len=c.seq_len,
@ -114,7 +113,7 @@ def train(model_G, criterion_G, optimizer_G, model_D, criterion_D, optimizer_D,
y_hat = model_G(c_G)
y_hat_sub = None
y_G_sub = None
y_hat_vis = y_hat # for visualization # FIXME! .clone().detach()
y_hat_vis = y_hat # for visualization
# PQMF formatting
if y_hat.shape[1] > 1:

View File

@ -8,12 +8,12 @@ import traceback
import numpy as np
import torch
from TTS.utils.arguments import parse_arguments, process_args
# DISTRIBUTED
from torch.nn.parallel import DistributedDataParallel as DDP_th
from torch.optim import Adam
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from TTS.utils.arguments import parse_arguments, process_args
from TTS.utils.audio import AudioProcessor
from TTS.utils.distribute import init_distributed
from TTS.utils.generic_utils import (KeepAverage, count_parameters,
@ -77,8 +77,8 @@ def format_test_data(data):
return m, x
def train(model, criterion, optimizer,
scheduler, scaler, ap, global_step, epoch):
def train(model, criterion, optimizer, scheduler, scaler, ap, global_step,
epoch):
data_loader = setup_loader(ap, is_val=False, verbose=(epoch == 0))
model.train()
epoch_time = 0
@ -92,7 +92,8 @@ def train(model, criterion, optimizer,
c_logger.print_train_start()
# setup noise schedule
noise_schedule = c['train_noise_schedule']
betas = np.linspace(noise_schedule['min_val'], noise_schedule['max_val'], noise_schedule['num_steps'])
betas = np.linspace(noise_schedule['min_val'], noise_schedule['max_val'],
noise_schedule['num_steps'])
if hasattr(model, 'module'):
model.module.compute_noise_level(betas)
else:
@ -193,7 +194,8 @@ def train(model, criterion, optimizer,
if global_step % c.save_step == 0:
if c.checkpoint:
# save model
save_checkpoint(model,
save_checkpoint(
model,
optimizer,
scheduler,
None,
@ -203,7 +205,8 @@ def train(model, criterion, optimizer,
epoch,
OUT_PATH,
model_losses=loss_dict,
scaler=scaler.state_dict() if c.mixed_precision else None)
scaler=scaler.state_dict() if c.mixed_precision else None
)
end_time = time.time()
@ -282,7 +285,9 @@ def evaluate(model, criterion, ap, global_step, epoch):
# setup noise schedule and inference
noise_schedule = c['test_noise_schedule']
betas = np.linspace(noise_schedule['min_val'], noise_schedule['max_val'], noise_schedule['num_steps'])
betas = np.linspace(noise_schedule['min_val'],
noise_schedule['max_val'],
noise_schedule['num_steps'])
if hasattr(model, 'module'):
model.module.compute_noise_level(betas)
# compute voice
@ -313,7 +318,8 @@ def main(args): # pylint: disable=redefined-outer-name
print(f" > Loading wavs from: {c.data_path}")
if c.feature_path is not None:
print(f" > Loading features from: {c.feature_path}")
eval_data, train_data = load_wav_feat_data(c.data_path, c.feature_path, c.eval_split_size)
eval_data, train_data = load_wav_feat_data(c.data_path, c.feature_path,
c.eval_split_size)
else:
eval_data, train_data = load_wav_data(c.data_path, c.eval_split_size)
@ -343,6 +349,10 @@ def main(args): # pylint: disable=redefined-outer-name
# setup criterion
criterion = torch.nn.L1Loss().cuda()
if use_cuda:
model.cuda()
criterion.cuda()
if args.restore_path:
checkpoint = torch.load(args.restore_path, map_location='cpu')
try:
@ -376,10 +386,6 @@ def main(args): # pylint: disable=redefined-outer-name
else:
args.restore_step = 0
if use_cuda:
model.cuda()
criterion.cuda()
# DISTRUBUTED
if num_gpus > 1:
model = DDP_th(model, device_ids=[args.rank])
@ -393,14 +399,13 @@ def main(args): # pylint: disable=redefined-outer-name
global_step = args.restore_step
for epoch in range(0, c.epochs):
c_logger.print_epoch_start(epoch, c.epochs)
_, global_step = train(model, criterion, optimizer,
scheduler, scaler, ap, global_step,
epoch)
eval_avg_loss_dict = evaluate(model, criterion, ap,
global_step, epoch)
_, global_step = train(model, criterion, optimizer, scheduler, scaler,
ap, global_step, epoch)
eval_avg_loss_dict = evaluate(model, criterion, ap, global_step, epoch)
c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
target_loss = eval_avg_loss_dict[c.target_loss]
best_loss = save_best_model(target_loss,
best_loss = save_best_model(
target_loss,
best_loss,
model,
optimizer,
@ -412,7 +417,8 @@ def main(args): # pylint: disable=redefined-outer-name
epoch,
OUT_PATH,
model_losses=eval_avg_loss_dict,
scaler=scaler.state_dict() if c.mixed_precision else None)
scaler=scaler.state_dict() if c.mixed_precision else None
)
if __name__ == '__main__':

View File

@ -178,7 +178,8 @@ def train(model, optimizer, criterion, scheduler, scaler, ap, global_step, epoch
if global_step % c.save_step == 0:
if c.checkpoint:
# save model
save_checkpoint(model,
save_checkpoint(
model,
optimizer,
scheduler,
None,
@ -204,14 +205,7 @@ def train(model, optimizer, criterion, scheduler, scaler, ap, global_step, epoch
c.batched,
c.target_samples,
c.overlap_samples,
# use_cuda
)
# sample_wav = model.generate(ground_mel,
# c.batched,
# c.target_samples,
# c.overlap_samples,
# use_cuda
# )
predict_mel = ap.melspectrogram(sample_wav)
# compute spectrograms
@ -300,7 +294,6 @@ def evaluate(model, criterion, ap, global_step, epoch):
c.batched,
c.target_samples,
c.overlap_samples,
# use_cuda
)
predict_mel = ap.melspectrogram(sample_wav)
@ -311,7 +304,8 @@ def evaluate(model, criterion, ap, global_step, epoch):
)
# compute spectrograms
figures = {"eval/ground_truth": plot_spectrogram(ground_mel.T),
figures = {
"eval/ground_truth": plot_spectrogram(ground_mel.T),
"eval/prediction": plot_spectrogram(predict_mel.T)
}
tb_logger.tb_eval_figures(global_step, figures)