Update train_ms.py

This commit is contained in:
Stardust·减
2023-08-03 18:29:47 +08:00
committed by GitHub
parent d32fa0e15b
commit 3ec81c2da1

View File

@@ -35,6 +35,10 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
from text.symbols import symbols from text.symbols import symbols
torch.backends.cudnn.benchmark = True torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision('medium')
global_step = 0 global_step = 0
@@ -109,12 +113,17 @@ def run(rank, n_gpus, hps):
pretrain_dir = None pretrain_dir = None
if pretrain_dir is None: if pretrain_dir is None:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, try:
optim_g, False) _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g,
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_g, skip_optimizer)
optim_d, False) _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d,
epoch_str = max(epoch_str, 1) optim_d, skip_optimizer)
global_step = (epoch_str - 1) * len(train_loader) epoch_str = max(epoch_str, 1)
global_step = (epoch_str - 1) * len(train_loader)
except Exception:
print("load old checkpoint failed...")
epoch_str = 1
global_step = 0
else: else:
_, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(pretrain_dir, "G_*.pth"), net_g, _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(pretrain_dir, "G_*.pth"), net_g,
optim_g, True) optim_g, True)
@@ -269,7 +278,6 @@ def evaluate(hps, generator, eval_loader, writer_eval):
print("Evaluating ...") print("Evaluating ...")
with torch.no_grad(): with torch.no_grad():
for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers, tone, language, bert) in enumerate(eval_loader): for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers, tone, language, bert) in enumerate(eval_loader):
print(111)
x, x_lengths = x.cuda(), x_lengths.cuda() x, x_lengths = x.cuda(), x_lengths.cuda()
spec, spec_lengths = spec.cuda(), spec_lengths.cuda() spec, spec_lengths = spec.cuda(), spec_lengths.cuda()
y, y_lengths = y.cuda(), y_lengths.cuda() y, y_lengths = y.cuda(), y_lengths.cuda()