tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Update train_ms.py

Browse Source
This commit is contained in:
Stardust·减
2023-08-23 18:38:51 +08:00
committed by GitHub
parent da97901831
commit c9b6a9695e
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
train_ms.py
View File

@@ -7,9 +7,8 @@ import torch
from torch import nn, optim
from torch.nn import functional as F
from torch.utils.data import DataLoader
#from torch.utils.tensorboard import SummaryWriter
import wandb
from wandb.tensorboard import WandbSummaryWriter as SummaryWriter
from torch.utils.tensorboard import SummaryWriter
#import wandb
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
@@ -64,8 +63,9 @@ def run(rank, n_gpus, hps):
logger = utils.get_logger(hps.model_dir)
logger.info(hps)
utils.check_git_hash(hps.model_dir)
wandb.init(project="Bert-VITS2", config=hps, sync_tensorboard=True)
#wandb.init(project="Bert-VITS2", config=hps, sync_tensorboard=True)
writer = SummaryWriter(log_dir=hps.model_dir)
#writer = wandb.summary.create_file_writer()
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
@@ -305,15 +305,15 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
"grad_norm_d": grad_norm_d, "grad_norm_g": grad_norm_g}
scalar_dict.update(
{"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl})
if net_dur_disc is not None:
scalar_dict.update({"loss/dur_disc/total": loss_dur_disc_all, "grad_norm_dur_disc": grad_norm_dur_disc})
#if net_dur_disc is not None:
#scalar_dict.update({"loss/dur_disc/total": loss_dur_disc_all, "grad_norm_dur_disc": grad_norm_dur_disc})
scalar_dict.update({"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)})
scalar_dict.update({"loss/d_r/{}".format(i): v for i, v in enumerate(losses_disc_r)})
scalar_dict.update({"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)})
if net_dur_disc is not None:
scalar_dict.update({"loss/dur_disc_r" : f"{losses_dur_disc_r}"})
scalar_dict.update({"loss/dur_disc_g" : f"{losses_dur_disc_g}"})
scalar_dict.update({"loss/dur_gen" : f"{loss_dur_gen}"})
#if net_dur_disc is not None:
#scalar_dict.update({"loss/dur_disc_r" : f"{losses_dur_disc_r}"})
#scalar_dict.update({"loss/dur_disc_g" : f"{losses_dur_disc_g}"})
#scalar_dict.update({"loss/dur_gen" : f"{loss_dur_gen}"})
image_dict = {
"slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),