Update train_ms.py

This commit is contained in:
Stardust·减
2023-08-23 18:19:45 +08:00
committed by GitHub
parent 81ef49f278
commit da97901831

View File

@@ -8,6 +8,7 @@ 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
from wandb.tensorboard import WandbSummaryWriter as SummaryWriter 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
@@ -63,6 +64,7 @@ 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)
writer = SummaryWriter(log_dir=hps.model_dir) writer = SummaryWriter(log_dir=hps.model_dir)
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval")) writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))