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

add stop loss

This commit is contained in:
Eren Golge 2018-04-30 05:47:14 -07:00
parent d4da61b78e
commit cc9bfe96af
1 changed files with 24 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
train.py
View File

@ -62,11 +62,12 @@ else:
print(" > Priority freq. is disabled.") print(" > Priority freq. is disabled.")
def train(model, criterion, data_loader, optimizer, epoch): def train(model, criterion, criterion_st, data_loader, optimizer, epoch):
model = model.train() model = model.train()
epoch_time = 0 epoch_time = 0
avg_linear_loss = 0 avg_linear_loss = 0
avg_mel_loss = 0 avg_mel_loss = 0
avg_stop_loss = 0
avg_attn_loss = 0 avg_attn_loss = 0
print(" | > Epoch {}/{}".format(epoch, c.epochs)) print(" | > Epoch {}/{}".format(epoch, c.epochs))
@ -108,18 +109,19 @@ def train(model, criterion, data_loader, optimizer, epoch):
mk = mk_decay(c.mk, c.epochs, epoch) mk = mk_decay(c.mk, c.epochs, epoch)
# forward pass # forward pass
mel_output, linear_output, alignments =\ mel_output, linear_output, alignments, stop_tokens =\
model.forward(text_input, mel_spec) model.forward(text_input, mel_spec)
# loss computation # loss computation
mel_loss = criterion(mel_output, mel_spec, mel_lengths) mel_loss = criterion(mel_output, mel_spec, mel_lengths)
linear_loss = criterion(linear_output, linear_spec, mel_lengths) linear_loss = criterion(linear_output, linear_spec, mel_lengths)
stop_loss = criterion_st(stop_tokens, stop_targets)
if c.priority_freq: if c.priority_freq:
linear_loss = 0.5 * linear_loss\ linear_loss = 0.5 * linear_loss\
+ 0.5 * criterion(linear_output[:, :, :n_priority_freq], + 0.5 * criterion(linear_output[:, :, :n_priority_freq],
linear_spec[:, :, :n_priority_freq], linear_spec[:, :, :n_priority_freq],
mel_lengths) mel_lengths)
loss = mel_loss + linear_loss loss = mel_loss + linear_loss + stop_loss
if c.mk > 0.0: if c.mk > 0.0:
attention_loss = criterion(alignments, M, mel_lengths) attention_loss = criterion(alignments, M, mel_lengths)
loss += mk * attention_loss loss += mk * attention_loss
@ -141,12 +143,14 @@ def train(model, criterion, data_loader, optimizer, epoch):
progbar_display['total_loss'] = loss.item() progbar_display['total_loss'] = loss.item()
progbar_display['linear_loss'] = linear_loss.item() progbar_display['linear_loss'] = linear_loss.item()
progbar_display['mel_loss'] = mel_loss.item() progbar_display['mel_loss'] = mel_loss.item()
progbar_display['stop_loss'] = stop_loss.item()
progbar_display['grad_norm'] = grad_norm.item() progbar_display['grad_norm'] = grad_norm.item()
# update # update
progbar.update(num_iter+1, values=list(progbar_display.items())) progbar.update(num_iter+1, values=list(progbar_display.items()))
avg_linear_loss += linear_loss.item() avg_linear_loss += linear_loss.item()
avg_mel_loss += mel_loss.item() avg_mel_loss += mel_loss.item()
avg_stop_loss += st_loss.item()
# Plot Training Iter Stats # Plot Training Iter Stats
tb.add_scalar('TrainIterLoss/TotalLoss', loss.item(), current_step) tb.add_scalar('TrainIterLoss/TotalLoss', loss.item(), current_step)
@ -193,11 +197,13 @@ def train(model, criterion, data_loader, optimizer, epoch):
avg_linear_loss /= (num_iter + 1) avg_linear_loss /= (num_iter + 1)
avg_mel_loss /= (num_iter + 1) avg_mel_loss /= (num_iter + 1)
avg_total_loss = avg_mel_loss + avg_linear_loss avg_stop_loss /= (num_iter + 1)
avg_total_loss = avg_mel_loss + avg_linear_loss + avg_stop_loss
# Plot Training Epoch Stats # Plot Training Epoch Stats
tb.add_scalar('TrainEpochLoss/TotalLoss', avg_total_loss, current_step) tb.add_scalar('TrainEpochLoss/TotalLoss', avg_total_loss, current_step)
tb.add_scalar('TrainEpochLoss/LinearLoss', avg_linear_loss, current_step) tb.add_scalar('TrainEpochLoss/LinearLoss', avg_linear_loss, current_step)
tb.add_scalar('TrainEpochLoss/StopLoss', avg_stop_loss, current_step)
tb.add_scalar('TrainEpochLoss/MelLoss', avg_mel_loss, current_step) tb.add_scalar('TrainEpochLoss/MelLoss', avg_mel_loss, current_step)
if c.mk > 0: if c.mk > 0:
avg_attn_loss /= (num_iter + 1) avg_attn_loss /= (num_iter + 1)
@ -208,11 +214,12 @@ def train(model, criterion, data_loader, optimizer, epoch):
return avg_linear_loss, current_step return avg_linear_loss, current_step
def evaluate(model, criterion, data_loader, current_step): def evaluate(model, criterion, criterion_st, data_loader, current_step):
model = model.eval() model = model.eval()
epoch_time = 0 epoch_time = 0
avg_linear_loss = 0 avg_linear_loss = 0
avg_mel_loss = 0 avg_mel_loss = 0
avg_stop_loss = 0
print("\n | > Validation") print("\n | > Validation")
progbar = Progbar(len(data_loader.dataset) / c.batch_size) progbar = Progbar(len(data_loader.dataset) / c.batch_size)
@ -236,18 +243,19 @@ def evaluate(model, criterion, data_loader, current_step):
linear_spec = linear_spec.cuda() linear_spec = linear_spec.cuda()
# forward pass # forward pass
mel_output, linear_output, alignments =\ mel_output, linear_output, alignments, stop_tokens =\
model.forward(text_input, mel_spec) model.forward(text_input, mel_spec)
# loss computation # loss computation
mel_loss = criterion(mel_output, mel_spec, mel_lengths) mel_loss = criterion(mel_output, mel_spec, mel_lengths)
linear_loss = criterion(linear_output, linear_spec, mel_lengths) linear_loss = criterion(linear_output, linear_spec, mel_lengths)
stop_loss = criterion_st(stop_tokens, stop_targets)
if c.priority_freq: if c.priority_freq:
linear_loss = 0.5 * linear_loss\ linear_loss = 0.5 * linear_loss\
+ 0.5 * criterion(linear_output[:, :, :n_priority_freq], + 0.5 * criterion(linear_output[:, :, :n_priority_freq],
linear_spec[:, :, :n_priority_freq], linear_spec[:, :, :n_priority_freq],
mel_lengths) mel_lengths)
loss = mel_loss + linear_loss loss = mel_loss + linear_loss + stop_loss
step_time = time.time() - start_time step_time = time.time() - start_time
epoch_time += step_time epoch_time += step_time
@ -256,11 +264,13 @@ def evaluate(model, criterion, data_loader, current_step):
progbar.update(num_iter+1, values=[('total_loss', loss.item()), progbar.update(num_iter+1, values=[('total_loss', loss.item()),
('linear_loss', ('linear_loss',
linear_loss.item()), linear_loss.item()),
('stop_loss', stop_loss.item()),
('mel_loss', mel_loss.item())]) ('mel_loss', mel_loss.item())])
sys.stdout.flush() sys.stdout.flush()
avg_linear_loss += linear_loss.item() avg_linear_loss += linear_loss.item()
avg_mel_loss += mel_loss.item() avg_mel_loss += mel_loss.item()
avg_stop_loss += stop_loss.item()
# Diagnostic visualizations # Diagnostic visualizations
idx = np.random.randint(mel_spec.shape[0]) idx = np.random.randint(mel_spec.shape[0])
@ -292,12 +302,14 @@ def evaluate(model, criterion, data_loader, current_step):
# compute average losses # compute average losses
avg_linear_loss /= (num_iter + 1) avg_linear_loss /= (num_iter + 1)
avg_mel_loss /= (num_iter + 1) avg_mel_loss /= (num_iter + 1)
avg_total_loss = avg_mel_loss + avg_linear_loss avg_stop_loss /= (num_iter + 1)
avg_total_loss = avg_mel_loss + avg_linear_loss + stop_loss
# Plot Learning Stats # Plot Learning Stats
tb.add_scalar('ValEpochLoss/TotalLoss', avg_total_loss, current_step) tb.add_scalar('ValEpochLoss/TotalLoss', avg_total_loss, current_step)
tb.add_scalar('ValEpochLoss/LinearLoss', avg_linear_loss, current_step) tb.add_scalar('ValEpochLoss/LinearLoss', avg_linear_loss, current_step)
tb.add_scalar('ValEpochLoss/MelLoss', avg_mel_loss, current_step) tb.add_scalar('ValEpochLoss/MelLoss', avg_mel_loss, current_step)
tb.add_scalar('ValEpochLoss/Stop_loss', avg_stop_loss, current_step)
return avg_linear_loss return avg_linear_loss
@ -355,8 +367,10 @@ def main(args):
if use_cuda: if use_cuda:
criterion = L1LossMasked().cuda() criterion = L1LossMasked().cuda()
criterion_st = nn.BCELoss().cuda()
else: else:
criterion = L1LossMasked() criterion = L1LossMasked()
criterion_st = nn.BCELoss()
if args.restore_path: if args.restore_path:
checkpoint = torch.load(args.restore_path) checkpoint = torch.load(args.restore_path)
@ -392,8 +406,8 @@ def main(args):
for epoch in range(0, c.epochs): for epoch in range(0, c.epochs):
train_loss, current_step = train( train_loss, current_step = train(
model, criterion, train_loader, optimizer, epoch) model, criterion, criterion_st, train_loader, optimizer, epoch)
val_loss = evaluate(model, criterion, val_loader, current_step) val_loss = evaluate(model, criterion, criterion_st, val_loader, current_step)
best_loss = save_best_model(model, optimizer, val_loss, best_loss = save_best_model(model, optimizer, val_loss,
best_loss, OUT_PATH, best_loss, OUT_PATH,
current_step, epoch) current_step, epoch)