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 import nn, optim
from torch.nn import functional as F from torch.nn import functional as F
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
#from torch.utils.tensorboard import SummaryWriter from torch.utils.tensorboard import SummaryWriter
import wandb #import wandb
from wandb.tensorboard import WandbSummaryWriter as SummaryWriter
import torch.multiprocessing as mp import torch.multiprocessing as mp
import torch.distributed as dist import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP 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 = utils.get_logger(hps.model_dir)
logger.info(hps) logger.info(hps)
utils.check_git_hash(hps.model_dir) 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 = SummaryWriter(log_dir=hps.model_dir)
#writer = wandb.summary.create_file_writer()
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval")) 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) 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} "grad_norm_d": grad_norm_d, "grad_norm_g": grad_norm_g}
scalar_dict.update( scalar_dict.update(
{"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl}) {"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: #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/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/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_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)}) scalar_dict.update({"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)})
if net_dur_disc is not None: #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_r" : f"{losses_dur_disc_r}"})
scalar_dict.update({"loss/dur_disc_g" : f"{losses_dur_disc_g}"}) #scalar_dict.update({"loss/dur_disc_g" : f"{losses_dur_disc_g}"})
scalar_dict.update({"loss/dur_gen" : f"{loss_dur_gen}"}) #scalar_dict.update({"loss/dur_gen" : f"{loss_dur_gen}"})
image_dict = { image_dict = {
"slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()), "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),