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

update train_align_tts.py for coqpit

This commit is contained in:
Eren Gölge 2021-05-06 15:53:32 +02:00
parent 51a7e06945
commit 7227e8f1d2
2 changed files with 575 additions and 464 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

304
TTS/bin/train_align_tts.py
View File

@ -23,62 +23,64 @@ from TTS.tts.utils.speakers import parse_speakers
from TTS.tts.utils.synthesis import synthesis from TTS.tts.utils.synthesis import synthesis
from TTS.tts.utils.text.symbols import make_symbols, phonemes, symbols from TTS.tts.utils.text.symbols import make_symbols, phonemes, symbols
from TTS.tts.utils.visual import plot_alignment, plot_spectrogram from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
from TTS.utils.arguments import parse_arguments, process_args from TTS.utils.arguments import init_training
from TTS.utils.audio import AudioProcessor from TTS.utils.audio import AudioProcessor
from TTS.utils.distribute import init_distributed, reduce_tensor from TTS.utils.distribute import init_distributed, reduce_tensor
from TTS.utils.generic_utils import KeepAverage, count_parameters, remove_experiment_folder, set_init_dict from TTS.utils.generic_utils import KeepAverage, count_parameters, remove_experiment_folder, set_init_dict
from TTS.utils.radam import RAdam from TTS.utils.radam import RAdam
from TTS.utils.training import NoamLR, setup_torch_training_env from TTS.utils.training import NoamLR, setup_torch_training_env
if __name__ == "__main__": use_cuda, num_gpus = setup_torch_training_env(True, False)
use_cuda, num_gpus = setup_torch_training_env(True, False) # torch.autograd.set_detect_anomaly(True)
# torch.autograd.set_detect_anomaly(True)
def setup_loader(ap, r, is_val=False, verbose=False):
if is_val and not c.run_eval: def setup_loader(ap, r, is_val=False, verbose=False):
if is_val and not config.run_eval:
loader = None loader = None
else: else:
dataset = MyDataset( dataset = MyDataset(
r, r,
c.text_cleaner, config.text_cleaner,
compute_linear_spec=False, compute_linear_spec=False,
meta_data=meta_data_eval if is_val else meta_data_train, meta_data=meta_data_eval if is_val else meta_data_train,
ap=ap, ap=ap,
tp=c.characters if "characters" in c.keys() else None, tp=config.characters,
add_blank=c["add_blank"] if "add_blank" in c.keys() else False, add_blank=config["add_blank"],
batch_group_size=0 if is_val else c.batch_group_size * c.batch_size, batch_group_size=0 if is_val else config.batch_group_size *
min_seq_len=c.min_seq_len, config.batch_size,
max_seq_len=c.max_seq_len, min_seq_len=config.min_seq_len,
phoneme_cache_path=c.phoneme_cache_path, max_seq_len=config.max_seq_len,
use_phonemes=c.use_phonemes, phoneme_cache_path=config.phoneme_cache_path,
phoneme_language=c.phoneme_language, use_phonemes=config.use_phonemes,
enable_eos_bos=c.enable_eos_bos_chars, phoneme_language=config.phoneme_language,
enable_eos_bos=config.enable_eos_bos_chars,
use_noise_augment=not is_val, use_noise_augment=not is_val,
verbose=verbose, verbose=verbose,
speaker_mapping=speaker_mapping speaker_mapping=speaker_mapping if config.use_speaker_embedding
if c.use_speaker_embedding and c.use_external_speaker_embedding_file and config.use_external_speaker_embedding_file else None,
else None,
) )
if c.use_phonemes and c.compute_input_seq_cache: if config.use_phonemes and config.compute_input_seq_cache:
# precompute phonemes to have a better estimate of sequence lengths. # precompute phonemes to have a better estimate of sequence lengths.
dataset.compute_input_seq(c.num_loader_workers) dataset.compute_input_seq(config.num_loader_workers)
dataset.sort_items() dataset.sort_items()
sampler = DistributedSampler(dataset) if num_gpus > 1 else None sampler = DistributedSampler(dataset) if num_gpus > 1 else None
loader = DataLoader( loader = DataLoader(
dataset, dataset,
batch_size=c.eval_batch_size if is_val else c.batch_size, batch_size=config.eval_batch_size if is_val else config.batch_size,
shuffle=False, shuffle=False,
collate_fn=dataset.collate_fn, collate_fn=dataset.collate_fn,
drop_last=False, drop_last=False,
sampler=sampler, sampler=sampler,
num_workers=c.num_val_loader_workers if is_val else c.num_loader_workers, num_workers=config.num_val_loader_workers
if is_val else config.num_loader_workers,
pin_memory=False, pin_memory=False,
) )
return loader return loader
def format_data(data):
def format_data(data):
# setup input data # setup input data
text_input = data[0] text_input = data[0]
text_lengths = data[1] text_lengths = data[1]
@ -89,13 +91,15 @@ if __name__ == "__main__":
avg_text_length = torch.mean(text_lengths.float()) avg_text_length = torch.mean(text_lengths.float())
avg_spec_length = torch.mean(mel_lengths.float()) avg_spec_length = torch.mean(mel_lengths.float())
if c.use_speaker_embedding: if config.use_speaker_embedding:
if c.use_external_speaker_embedding_file: if config.use_external_speaker_embedding_file:
# return precomputed embedding vector # return precomputed embedding vector
speaker_c = data[8] speaker_c = data[8]
else: else:
# return speaker_id to be used by an embedding layer # return speaker_id to be used by an embedding layer
speaker_c = [speaker_mapping[speaker_name] for speaker_name in speaker_names] speaker_c = [
speaker_mapping[speaker_name] for speaker_name in speaker_names
]
speaker_c = torch.LongTensor(speaker_c) speaker_c = torch.LongTensor(speaker_c)
else: else:
speaker_c = None speaker_c = None
@ -109,18 +113,21 @@ if __name__ == "__main__":
speaker_c = speaker_c.cuda(non_blocking=True) speaker_c = speaker_c.cuda(non_blocking=True)
return text_input, text_lengths, mel_input, mel_lengths, speaker_c, avg_text_length, avg_spec_length, item_idx return text_input, text_lengths, mel_input, mel_lengths, speaker_c, avg_text_length, avg_spec_length, item_idx
def train(data_loader, model, criterion, optimizer, scheduler, ap, global_step, epoch, training_phase):
def train(data_loader, model, criterion, optimizer, scheduler, ap, global_step,
epoch, training_phase):
model.train() model.train()
epoch_time = 0 epoch_time = 0
keep_avg = KeepAverage() keep_avg = KeepAverage()
if use_cuda: if use_cuda:
batch_n_iter = int(len(data_loader.dataset) / (c.batch_size * num_gpus)) batch_n_iter = int(
len(data_loader.dataset) / (config.batch_size * num_gpus))
else: else:
batch_n_iter = int(len(data_loader.dataset) / c.batch_size) batch_n_iter = int(len(data_loader.dataset) / config.batch_size)
end_time = time.time() end_time = time.time()
c_logger.print_train_start() c_logger.print_train_start()
scaler = torch.cuda.amp.GradScaler() if c.mixed_precision else None scaler = torch.cuda.amp.GradScaler() if config.mixed_precision else None
for num_iter, data in enumerate(data_loader): for num_iter, data in enumerate(data_loader):
start_time = time.time() start_time = time.time()
@ -142,10 +149,14 @@ if __name__ == "__main__":
optimizer.zero_grad() optimizer.zero_grad()
# forward pass model # forward pass model
with torch.cuda.amp.autocast(enabled=c.mixed_precision): with torch.cuda.amp.autocast(enabled=config.mixed_precision):
decoder_output, dur_output, dur_mas_output, alignments, _, _, logp = model.forward( decoder_output, dur_output, dur_mas_output, alignments, _, _, logp = model.forward(
text_input, text_lengths, mel_targets, mel_lengths, g=speaker_c, phase=training_phase text_input,
) text_lengths,
mel_targets,
mel_lengths,
g=speaker_c,
phase=training_phase)
# compute loss # compute loss
loss_dict = criterion( loss_dict = criterion(
@ -161,19 +172,21 @@ if __name__ == "__main__":
) )
# backward pass with loss scaling # backward pass with loss scaling
if c.mixed_precision: if config.mixed_precision:
scaler.scale(loss_dict["loss"]).backward() scaler.scale(loss_dict["loss"]).backward()
scaler.unscale_(optimizer) scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), c.grad_clip) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
config.grad_clip)
scaler.step(optimizer) scaler.step(optimizer)
scaler.update() scaler.update()
else: else:
loss_dict["loss"].backward() loss_dict["loss"].backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), c.grad_clip) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
config.grad_clip)
optimizer.step() optimizer.step()
# setup lr # setup lr
if c.noam_schedule: if config.noam_schedule:
scheduler.step() scheduler.step()
# current_lr # current_lr
@ -188,9 +201,12 @@ if __name__ == "__main__":
# aggregate losses from processes # aggregate losses from processes
if num_gpus > 1: if num_gpus > 1:
loss_dict["loss_l1"] = reduce_tensor(loss_dict["loss_l1"].data, num_gpus) loss_dict["loss_l1"] = reduce_tensor(loss_dict["loss_l1"].data,
loss_dict["loss_ssim"] = reduce_tensor(loss_dict["loss_ssim"].data, num_gpus) num_gpus)
loss_dict["loss_dur"] = reduce_tensor(loss_dict["loss_dur"].data, num_gpus) loss_dict["loss_ssim"] = reduce_tensor(loss_dict["loss_ssim"].data,
num_gpus)
loss_dict["loss_dur"] = reduce_tensor(loss_dict["loss_dur"].data,
num_gpus)
loss_dict["loss"] = reduce_tensor(loss_dict["loss"].data, num_gpus) loss_dict["loss"] = reduce_tensor(loss_dict["loss"].data, num_gpus)
# detach loss values # detach loss values
@ -211,7 +227,7 @@ if __name__ == "__main__":
keep_avg.update_values(update_train_values) keep_avg.update_values(update_train_values)
# print training progress # print training progress
if global_step % c.print_step == 0: if global_step % config.print_step == 0:
log_dict = { log_dict = {
"avg_spec_length": [avg_spec_length, 1], # value, precision "avg_spec_length": [avg_spec_length, 1], # value, precision
"avg_text_length": [avg_text_length, 1], "avg_text_length": [avg_text_length, 1],
@ -219,18 +235,23 @@ if __name__ == "__main__":
"loader_time": [loader_time, 2], "loader_time": [loader_time, 2],
"current_lr": current_lr, "current_lr": current_lr,
} }
c_logger.print_train_step(batch_n_iter, num_iter, global_step, log_dict, loss_dict, keep_avg.avg_values) c_logger.print_train_step(batch_n_iter, num_iter, global_step,
log_dict, loss_dict, keep_avg.avg_values)
if args.rank == 0: if args.rank == 0:
# Plot Training Iter Stats # Plot Training Iter Stats
# reduce TB load # reduce TB load
if global_step % c.tb_plot_step == 0: if global_step % config.tb_plot_step == 0:
iter_stats = {"lr": current_lr, "grad_norm": grad_norm, "step_time": step_time} iter_stats = {
"lr": current_lr,
"grad_norm": grad_norm,
"step_time": step_time
}
iter_stats.update(loss_dict) iter_stats.update(loss_dict)
tb_logger.tb_train_iter_stats(global_step, iter_stats) tb_logger.tb_train_iter_stats(global_step, iter_stats)
if global_step % c.save_step == 0: if global_step % config.save_step == 0:
if c.checkpoint: if config.checkpoint:
# save model # save model
save_checkpoint( save_checkpoint(
model, model,
@ -249,7 +270,8 @@ if __name__ == "__main__":
# Diagnostic visualizations # Diagnostic visualizations
if decoder_output is not None: if decoder_output is not None:
idx = np.random.randint(mel_targets.shape[0]) idx = np.random.randint(mel_targets.shape[0])
pred_spec = decoder_output[idx].detach().data.cpu().numpy().T pred_spec = decoder_output[idx].detach().data.cpu().numpy(
).T
gt_spec = mel_targets[idx].data.cpu().numpy().T gt_spec = mel_targets[idx].data.cpu().numpy().T
align_img = alignments[idx].data.cpu() align_img = alignments[idx].data.cpu()
@ -263,7 +285,9 @@ if __name__ == "__main__":
# Sample audio # Sample audio
train_audio = ap.inv_melspectrogram(pred_spec.T) train_audio = ap.inv_melspectrogram(pred_spec.T)
tb_logger.tb_train_audios(global_step, {"TrainAudio": train_audio}, c.audio["sample_rate"]) tb_logger.tb_train_audios(global_step,
{"TrainAudio": train_audio},
config.audio["sample_rate"])
end_time = time.time() end_time = time.time()
# print epoch stats # print epoch stats
@ -274,12 +298,14 @@ if __name__ == "__main__":
epoch_stats = {"epoch_time": epoch_time} epoch_stats = {"epoch_time": epoch_time}
epoch_stats.update(keep_avg.avg_values) epoch_stats.update(keep_avg.avg_values)
tb_logger.tb_train_epoch_stats(global_step, epoch_stats) tb_logger.tb_train_epoch_stats(global_step, epoch_stats)
if c.tb_model_param_stats: if config.tb_model_param_stats:
tb_logger.tb_model_weights(model, global_step) tb_logger.tb_model_weights(model, global_step)
return keep_avg.avg_values, global_step return keep_avg.avg_values, global_step
@torch.no_grad()
def evaluate(data_loader, model, criterion, ap, global_step, epoch, training_phase): @torch.no_grad()
def evaluate(data_loader, model, criterion, ap, global_step, epoch,
training_phase):
model.eval() model.eval()
epoch_time = 0 epoch_time = 0
keep_avg = KeepAverage() keep_avg = KeepAverage()
@ -289,13 +315,18 @@ if __name__ == "__main__":
start_time = time.time() start_time = time.time()
# format data # format data
text_input, text_lengths, mel_targets, mel_lengths, speaker_c, _, _, _ = format_data(data) text_input, text_lengths, mel_targets, mel_lengths, speaker_c, _, _, _ = format_data(
data)
# forward pass model # forward pass model
with torch.cuda.amp.autocast(enabled=c.mixed_precision): with torch.cuda.amp.autocast(enabled=config.mixed_precision):
decoder_output, dur_output, dur_mas_output, alignments, _, _, logp = model.forward( decoder_output, dur_output, dur_mas_output, alignments, _, _, logp = model.forward(
text_input, text_lengths, mel_targets, mel_lengths, g=speaker_c, phase=training_phase text_input,
) text_lengths,
mel_targets,
mel_lengths,
g=speaker_c,
phase=training_phase)
# compute loss # compute loss
loss_dict = criterion( loss_dict = criterion(
@ -320,10 +351,14 @@ if __name__ == "__main__":
# aggregate losses from processes # aggregate losses from processes
if num_gpus > 1: if num_gpus > 1:
loss_dict["loss_l1"] = reduce_tensor(loss_dict["loss_l1"].data, num_gpus) loss_dict["loss_l1"] = reduce_tensor(loss_dict["loss_l1"].data,
loss_dict["loss_ssim"] = reduce_tensor(loss_dict["loss_ssim"].data, num_gpus) num_gpus)
loss_dict["loss_dur"] = reduce_tensor(loss_dict["loss_dur"].data, num_gpus) loss_dict["loss_ssim"] = reduce_tensor(
loss_dict["loss"] = reduce_tensor(loss_dict["loss"].data, num_gpus) loss_dict["loss_ssim"].data, num_gpus)
loss_dict["loss_dur"] = reduce_tensor(
loss_dict["loss_dur"].data, num_gpus)
loss_dict["loss"] = reduce_tensor(loss_dict["loss"].data,
num_gpus)
# detach loss values # detach loss values
loss_dict_new = dict() loss_dict_new = dict()
@ -340,8 +375,9 @@ if __name__ == "__main__":
update_train_values["avg_" + key] = value update_train_values["avg_" + key] = value
keep_avg.update_values(update_train_values) keep_avg.update_values(update_train_values)
if c.print_eval: if config.print_eval:
c_logger.print_eval_step(num_iter, loss_dict, keep_avg.avg_values) c_logger.print_eval_step(num_iter, loss_dict,
keep_avg.avg_values)
if args.rank == 0: if args.rank == 0:
# Diagnostic visualizations # Diagnostic visualizations
@ -351,21 +387,27 @@ if __name__ == "__main__":
align_img = alignments[idx].data.cpu() align_img = alignments[idx].data.cpu()
eval_figures = { eval_figures = {
"prediction": plot_spectrogram(pred_spec, ap, output_fig=False), "prediction": plot_spectrogram(pred_spec, ap,
"ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False), output_fig=False),
"ground_truth": plot_spectrogram(gt_spec, ap,
output_fig=False),
"alignment": plot_alignment(align_img, output_fig=False), "alignment": plot_alignment(align_img, output_fig=False),
} }
# Sample audio # Sample audio
eval_audio = ap.inv_melspectrogram(pred_spec.T) eval_audio = ap.inv_melspectrogram(pred_spec.T)
tb_logger.tb_eval_audios(global_step, {"ValAudio": eval_audio}, c.audio["sample_rate"]) tb_logger.tb_eval_audios(global_step, {"ValAudio": eval_audio},
config.audio["sample_rate"])
# Plot Validation Stats # Plot Validation Stats
tb_logger.tb_eval_stats(global_step, keep_avg.avg_values) tb_logger.tb_eval_stats(global_step, keep_avg.avg_values)
tb_logger.tb_eval_figures(global_step, eval_figures) tb_logger.tb_eval_figures(global_step, eval_figures)
if args.rank == 0 and epoch >= c.test_delay_epochs: if args.rank == 0 and epoch >= config.test_delay_epochs:
if c.test_sentences_file is None: if config.test_sentences_file:
with open(config.test_sentences_file, "r") as f:
test_sentences = [s.strip() for s in f.readlines()]
else:
test_sentences = [ test_sentences = [
"It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.", "It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.",
"Be a voice, not an echo.", "Be a voice, not an echo.",
@ -373,19 +415,16 @@ if __name__ == "__main__":
"This cake is great. It's so delicious and moist.", "This cake is great. It's so delicious and moist.",
"Prior to November 22, 1963.", "Prior to November 22, 1963.",
] ]
else:
with open(c.test_sentences_file, "r") as f:
test_sentences = [s.strip() for s in f.readlines()]
# test sentences # test sentences
test_audios = {} test_audios = {}
test_figures = {} test_figures = {}
print(" | > Synthesizing test sentences") print(" | > Synthesizing test sentences")
if c.use_speaker_embedding: if config.use_speaker_embedding:
if c.use_external_speaker_embedding_file: if config.use_external_speaker_embedding_file:
speaker_embedding = speaker_mapping[ speaker_embedding = speaker_mapping[list(
list(speaker_mapping.keys())[randrange(len(speaker_mapping) - 1)] speaker_mapping.keys())[randrange(
]["embedding"] len(speaker_mapping) - 1)]]["embedding"]
speaker_id = None speaker_id = None
else: else:
speaker_id = 0 speaker_id = 0
@ -394,70 +433,79 @@ if __name__ == "__main__":
speaker_id = None speaker_id = None
speaker_embedding = None speaker_embedding = None
style_wav = c.get("style_wav_for_test")
for idx, test_sentence in enumerate(test_sentences): for idx, test_sentence in enumerate(test_sentences):
try: try:
wav, alignment, _, postnet_output, _, _ = synthesis( wav, alignment, _, postnet_output, _, _ = synthesis(
model, model,
test_sentence, test_sentence,
c, config,
use_cuda, use_cuda,
ap, ap,
speaker_id=speaker_id, speaker_id=speaker_id,
speaker_embedding=speaker_embedding, speaker_embedding=speaker_embedding,
style_wav=style_wav, style_wav=None,
truncated=False, truncated=False,
enable_eos_bos_chars=c.enable_eos_bos_chars, # pylint: disable=unused-argument enable_eos_bos_chars=config.enable_eos_bos_chars, # pylint: disable=unused-argument
use_griffin_lim=True, use_griffin_lim=True,
do_trim_silence=False, do_trim_silence=False,
) )
file_path = os.path.join(AUDIO_PATH, str(global_step)) file_path = os.path.join(AUDIO_PATH, str(global_step))
os.makedirs(file_path, exist_ok=True) os.makedirs(file_path, exist_ok=True)
file_path = os.path.join(file_path, "TestSentence_{}.wav".format(idx)) file_path = os.path.join(file_path,
"TestSentence_{}.wav".format(idx))
ap.save_wav(wav, file_path) ap.save_wav(wav, file_path)
test_audios["{}-audio".format(idx)] = wav test_audios["{}-audio".format(idx)] = wav
test_figures["{}-prediction".format(idx)] = plot_spectrogram(postnet_output, ap) test_figures["{}-prediction".format(idx)] = plot_spectrogram(
test_figures["{}-alignment".format(idx)] = plot_alignment(alignment) postnet_output, ap)
test_figures["{}-alignment".format(idx)] = plot_alignment(
alignment)
except: # pylint: disable=bare-except except: # pylint: disable=bare-except
print(" !! Error creating Test Sentence -", idx) print(" !! Error creating Test Sentence -", idx)
traceback.print_exc() traceback.print_exc()
tb_logger.tb_test_audios(global_step, test_audios, c.audio["sample_rate"]) tb_logger.tb_test_audios(global_step, test_audios,
config.audio["sample_rate"])
tb_logger.tb_test_figures(global_step, test_figures) tb_logger.tb_test_figures(global_step, test_figures)
return keep_avg.avg_values return keep_avg.avg_values
def main(args): # pylint: disable=redefined-outer-name
def main(args): # pylint: disable=redefined-outer-name
# pylint: disable=global-variable-undefined # pylint: disable=global-variable-undefined
global meta_data_train, meta_data_eval, symbols, phonemes, model_characters, speaker_mapping global meta_data_train, meta_data_eval, symbols, phonemes, model_characters, speaker_mapping
# Audio processor # Audio processor
ap = AudioProcessor(**c.audio) ap = AudioProcessor(**config.audio.to_dict())
if "characters" in c.keys(): if config.has("characters") and config.characters:
symbols, phonemes = make_symbols(**c.characters) symbols, phonemes = make_symbols(**config.characters.to_dict())
# DISTRUBUTED # DISTRUBUTED
if num_gpus > 1: if num_gpus > 1:
init_distributed(args.rank, num_gpus, args.group_id, c.distributed["backend"], c.distributed["url"]) init_distributed(args.rank, num_gpus, args.group_id,
config.distributed["backend"],
config.distributed["url"])
# set model characters # set model characters
model_characters = phonemes if c.use_phonemes else symbols model_characters = phonemes if config.use_phonemes else symbols
num_chars = len(model_characters) num_chars = len(model_characters)
# load data instances # load data instances
meta_data_train, meta_data_eval = load_meta_data(c.datasets, eval_split=True) meta_data_train, meta_data_eval = load_meta_data(config.datasets,
eval_split=True)
# set the portion of the data used for training if set in config.json
if "train_portion" in c.keys():
meta_data_train = meta_data_train[: int(len(meta_data_train) * c.train_portion)]
if "eval_portion" in c.keys():
meta_data_eval = meta_data_eval[: int(len(meta_data_eval) * c.eval_portion)]
# parse speakers # parse speakers
num_speakers, speaker_embedding_dim, speaker_mapping = parse_speakers(c, args, meta_data_train, OUT_PATH) num_speakers, speaker_embedding_dim, speaker_mapping = parse_speakers(
config, args, meta_data_train, OUT_PATH)
# setup model # setup model
model = setup_model(num_chars, num_speakers, c, speaker_embedding_dim=speaker_embedding_dim) model = setup_model(num_chars,
optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0, betas=(0.9, 0.98), eps=1e-9) num_speakers,
criterion = AlignTTSLoss(c) config,
speaker_embedding_dim=speaker_embedding_dim)
optimizer = RAdam(model.parameters(),
lr=config.lr,
weight_decay=0,
betas=(0.9, 0.98),
eps=1e-9)
criterion = AlignTTSLoss(config)
if args.restore_path: if args.restore_path:
print(f" > Restoring from {os.path.basename(args.restore_path)} ...") print(f" > Restoring from {os.path.basename(args.restore_path)} ...")
@ -466,19 +514,20 @@ if __name__ == "__main__":
# TODO: fix optimizer init, model.cuda() needs to be called before # TODO: fix optimizer init, model.cuda() needs to be called before
# optimizer restore # optimizer restore
optimizer.load_state_dict(checkpoint["optimizer"]) optimizer.load_state_dict(checkpoint["optimizer"])
if c.reinit_layers: if config.reinit_layers:
raise RuntimeError raise RuntimeError
model.load_state_dict(checkpoint["model"]) model.load_state_dict(checkpoint["model"])
except: # pylint: disable=bare-except except: # pylint: disable=bare-except
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"], config)
model.load_state_dict(model_dict) model.load_state_dict(model_dict)
del model_dict del model_dict
for group in optimizer.param_groups: for group in optimizer.param_groups:
group["initial_lr"] = c.lr group["initial_lr"] = config.lr
print(" > Model restored from step %d" % checkpoint["step"], flush=True) print(" > Model restored from step %d" % checkpoint["step"],
flush=True)
args.restore_step = checkpoint["step"] args.restore_step = checkpoint["step"]
else: else:
args.restore_step = 0 args.restore_step = 0
@ -491,8 +540,10 @@ if __name__ == "__main__":
if num_gpus > 1: if num_gpus > 1:
model = DDP_th(model, device_ids=[args.rank]) model = DDP_th(model, device_ids=[args.rank])
if c.noam_schedule: if config.noam_schedule:
scheduler = NoamLR(optimizer, warmup_steps=c.warmup_steps, last_epoch=args.restore_step - 1) scheduler = NoamLR(optimizer,
warmup_steps=config.warmup_steps,
last_epoch=args.restore_step - 1)
else: else:
scheduler = None scheduler = None
@ -503,11 +554,13 @@ if __name__ == "__main__":
best_loss = float("inf") best_loss = float("inf")
print(" > Starting with inf best loss.") print(" > Starting with inf best loss.")
else: else:
print(" > Restoring best loss from " f"{os.path.basename(args.best_path)} ...") print(" > Restoring best loss from "
best_loss = torch.load(args.best_path, map_location="cpu")["model_loss"] 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 loaded last best loss {best_loss}.")
keep_all_best = c.get("keep_all_best", False) keep_all_best = config.keep_all_best
keep_after = c.get("keep_after", 10000) # void if keep_all_best False keep_after = config.keep_after # void if keep_all_best False
# define dataloaders # define dataloaders
train_loader = setup_loader(ap, 1, is_val=False, verbose=True) train_loader = setup_loader(ap, 1, is_val=False, verbose=True)
@ -517,29 +570,32 @@ if __name__ == "__main__":
def set_phase(): def set_phase():
"""Set AlignTTS training phase""" """Set AlignTTS training phase"""
if isinstance(c.phase_start_steps, list): if isinstance(config.phase_start_steps, list):
vals = [i < global_step for i in c.phase_start_steps] vals = [i < global_step for i in config.phase_start_steps]
if not True in vals: if not True in vals:
phase = 0 phase = 0
else: else:
phase = ( phase = (
len(c.phase_start_steps) - [i < global_step for i in c.phase_start_steps][::-1].index(True) - 1 len(config.phase_start_steps) -
) [i < global_step
for i in config.phase_start_steps][::-1].index(True) - 1)
else: else:
phase = None phase = None
return phase return phase
for epoch in range(0, c.epochs): for epoch in range(0, config.epochs):
cur_phase = set_phase() cur_phase = set_phase()
print(f"\n > Current AlignTTS phase: {cur_phase}") print(f"\n > Current AlignTTS phase: {cur_phase}")
c_logger.print_epoch_start(epoch, c.epochs) c_logger.print_epoch_start(epoch, config.epochs)
train_avg_loss_dict, global_step = train( train_avg_loss_dict, global_step = train(train_loader, model,
train_loader, model, criterion, optimizer, scheduler, ap, global_step, epoch, cur_phase criterion, optimizer,
) scheduler, ap, global_step,
eval_avg_loss_dict = evaluate(eval_loader, model, criterion, ap, global_step, epoch, cur_phase) epoch, cur_phase)
eval_avg_loss_dict = evaluate(eval_loader, model, criterion, ap,
global_step, epoch, cur_phase)
c_logger.print_epoch_end(epoch, eval_avg_loss_dict) c_logger.print_epoch_end(epoch, eval_avg_loss_dict)
target_loss = train_avg_loss_dict["avg_loss"] target_loss = train_avg_loss_dict["avg_loss"]
if c.run_eval: if config.run_eval:
target_loss = eval_avg_loss_dict["avg_loss"] target_loss = eval_avg_loss_dict["avg_loss"]
best_loss = save_best_model( best_loss = save_best_model(
target_loss, target_loss,
@ -555,8 +611,10 @@ if __name__ == "__main__":
keep_after=keep_after, keep_after=keep_after,
) )
args = parse_arguments(sys.argv)
c, OUT_PATH, AUDIO_PATH, c_logger, tb_logger = process_args(args, model_class="tts") if __name__ == "__main__":
args, config, OUT_PATH, AUDIO_PATH, c_logger, tb_logger = init_training(
sys.argv)
try: try:
main(args) main(args)

53
TTS/tts/configs/align_tts_config.py Normal file
View File

@ -0,0 +1,53 @@
from dataclasses import dataclass, field
from .shared_configs import BaseTTSConfig
@dataclass
class AlignTTSConfig(BaseTTSConfig):
"""Defines parameters for AlignTTS model."""
model: str = "align_tts"
# model specific params
positional_encoding: bool = True
hidden_channels_dp: int = 256
hidden_channels: int = 256
encoder_type: str = "fftransformer"
encoder_params: dict = field(
default_factory=lambda: {
"hidden_channels_ffn": 1024,
"num_heads": 2,
"num_layers": 6,
"dropout_p": 0.1
})
decoder_type: str = "fftransformer"
decoder_params: dict = field(
default_factory=lambda: {
"hidden_channels_ffn": 1024,
"num_heads": 2,
"num_layers": 6,
"dropout_p": 0.1
})
phase_start_steps: list = None
ssim_alpha: float = 1.0
spec_loss_alpha: float = 1.0
dur_loss_alpha: float = 1.0
mdn_alpha: float = 1.0
# multi-speaker settings
use_speaker_embedding: bool = False
use_external_speaker_embedding_file: bool = False
external_speaker_embedding_file: str = False
# optimizer parameters
noam_schedule: bool = False
warmup_steps: int = 4000
lr: float = 1e-4
wd: float = 1e-6
grad_clip: float = 5.0
# overrides
min_seq_len: int = 13
max_seq_len: int = 200
r: int = 1