import argparse import os import json import torch import yaml import torch.nn as nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm from utils.model import get_model, get_vocoder, get_param_num from utils.tools import get_configs_of, to_device, log, synth_one_sample from model import Tacotron2Loss from dataset import Dataset import logging numba_logger = logging.getLogger('numba') numba_logger.setLevel(logging.WARNING) from evaluate import evaluate device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def main(args, configs): print("Prepare training ...") preprocess_config, model_config, train_config = configs # Get dataset dataset = Dataset( "train.txt", preprocess_config, model_config, train_config, sort=True, drop_last=True ) batch_size = train_config["optimizer"]["batch_size"] group_size = 4 # Set this larger than 1 to enable sorting in Dataset assert batch_size * group_size < len(dataset) loader = DataLoader( dataset, batch_size=batch_size * group_size, shuffle=True, collate_fn=dataset.collate_fn, ) # Prepare model model, optimizer = get_model(args, configs, device, train=True) if len(args.pretrained): print(f"Warm starting from checkpoint {args.pretrained}") ckpt = torch.load(args.pretrained) model.load_state_dict(ckpt["model"]) if args.checkpoint_interval > 1: print(f"Overriding checkpoint interval to {args.checkpoint_interval}") train_config["step"]["save_step"] = args.checkpoint_interval model = nn.DataParallel(model) num_param = get_param_num(model) Loss = Tacotron2Loss(preprocess_config, model_config, train_config).to(device) print("Number of Tacotron2 Parameters:", num_param) # Load vocoder vocoder = get_vocoder(model_config, device) # Init logger if len(args.out_path): train_config["path"]["log_path"] = os.path.join(args.out_path, "logs") train_config["path"]["ckpt_path"] = os.path.join(args.out_path, "checkpoints") train_config["path"]["result_path"] = os.path.join(args.out_path, "val") for p in train_config["path"].values(): os.makedirs(p, exist_ok=True) train_log_path = os.path.join(train_config["path"]["log_path"], "train") val_log_path = os.path.join(train_config["path"]["log_path"], "val") os.makedirs(train_log_path, exist_ok=True) os.makedirs(val_log_path, exist_ok=True) train_logger = SummaryWriter(train_log_path) val_logger = SummaryWriter(val_log_path) # Training step = args.restore_step + 1 epoch = 1 grad_acc_step = train_config["optimizer"]["grad_acc_step"] grad_clip_thresh = train_config["optimizer"]["grad_clip_thresh"] total_step = train_config["step"]["total_step"] log_step = train_config["step"]["log_step"] save_step = train_config["step"]["save_step"] synth_step = train_config["step"]["synth_step"] val_step = train_config["step"]["val_step"] normalize = preprocess_config["preprocessing"]["mel"]["normalize"] with open( os.path.join(preprocess_config["path"]["preprocessed_path"], "stats.json") ) as f: stats = json.load(f) mel_stats = stats["mel"] outer_bar = tqdm(total=total_step, desc="Training", position=0) outer_bar.n = args.restore_step outer_bar.update() while True: inner_bar = tqdm(total=len(loader), desc="Epoch {}".format(epoch), position=1) for batchs in loader: for batch in batchs: batch = to_device(batch, device, mel_stats if normalize else None) # Forward output = model(*(batch[2:])) # Cal Loss losses = Loss(batch, output) total_loss = losses[0] # Backward total_loss = total_loss / grad_acc_step total_loss.backward() if step % grad_acc_step == 0: # Clipping gradients to avoid gradient explosion grad_norm = nn.utils.clip_grad_norm_(model.parameters(), grad_clip_thresh) # Update weights optimizer.step_and_update_lr() optimizer.zero_grad() if step % log_step == 0: losses = [l.item() for l in losses] message1 = "Step {}/{}, ".format(step, total_step) message2 = "Total Loss: {:.4f}, Mel Loss: {:.4f}, Gate Loss: {:.4f}, Forward Sum Loss: {:.4f}, Guided Attention Loss: {:.4f}".format( *losses ) with open(os.path.join(train_log_path, "log.txt"), "a") as f: f.write(message1 + message2 + "\n") outer_bar.write(message1 + message2) log(train_logger, step, losses=losses, grad_norm=grad_norm) if step % synth_step == 0: fig, gate_fig, wav_reconstruction, wav_prediction, tag = synth_one_sample( batch, output, vocoder, mel_stats, model_config, preprocess_config, step, ) log( train_logger, fig=fig, tag="Training/step_{}_{}".format(step, tag), ) log( train_logger, step=step, fig=gate_fig, tag="Gates/training", ) sampling_rate = preprocess_config["preprocessing"]["audio"][ "sampling_rate" ] log( train_logger, audio=wav_reconstruction, sampling_rate=sampling_rate, tag="Training/step_{}_{}_reconstructed".format(step, tag), ) log( train_logger, audio=wav_prediction, sampling_rate=sampling_rate, tag="Training/step_{}_{}_synthesized".format(step, tag), ) if step % val_step == 0: model.eval() message = evaluate(model, step, configs, mel_stats, val_logger, vocoder, len(losses)) with open(os.path.join(val_log_path, "log.txt"), "a") as f: f.write(message + "\n") outer_bar.write(message) model.train() if step % save_step == 0: torch.save( { "model": model.module.state_dict(), "optimizer": optimizer._optimizer.state_dict(), }, os.path.join( train_config["path"]["ckpt_path"], "{}.pth.tar".format(step), ), ) if step == total_step: quit() step += 1 outer_bar.update(1) inner_bar.update(1) epoch += 1 if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--restore_step", type=int, default=0) parser.add_argument( "--dataset", type=str, required=True, help="name of dataset", ) parser.add_argument( "--pretrained", type=str, required=False, help="Path to pretrained checkpoint file to finetune from", default="", ) parser.add_argument( "--out-path", type=str, required=False, help="Path override for train", default="", ) parser.add_argument( "--checkpoint-interval", type=int, required=False, help="Checkpoint interval override for train", default=0, ) args = parser.parse_args() # Read Config configs = get_configs_of(args.dataset) main(args, configs)