from accelerate.logging import get_logger import logging from torch.utils.tensorboard import SummaryWriter from accelerate import DistributedDataParallelKwargs from accelerate.utils import LoggerType from accelerate import Accelerator import time from optimizers import build_optimizer from losses import * from utils import * from meldataset import build_dataloader from models import * import librosa import torchaudio import torch.nn.functional as F from torch import nn from munch import Munch import random import os import os.path as osp import re import sys import yaml import shutil import numpy as np import torch import click import warnings warnings.simplefilter('ignore') # load packages torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision("high") logger = get_logger(__name__, log_level="DEBUG") @click.command() @click.option('-p', '--config_path', default='Configs/config.yml', type=str) def main(config_path): config = yaml.safe_load(open(config_path)) log_dir = config['log_dir'] if not osp.exists(log_dir): os.makedirs(log_dir, exist_ok=True) shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path))) ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( project_dir=log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs]) if accelerator.is_main_process: writer = SummaryWriter(log_dir + "/tensorboard") # write logs file_handler = logging.FileHandler(osp.join(log_dir, 'train.log')) file_handler.setLevel(logging.DEBUG) file_handler.setFormatter(logging.Formatter( '%(levelname)s:%(asctime)s: %(message)s')) logger.logger.addHandler(file_handler) batch_size = config.get('batch_size', 10) device = accelerator.device epochs = config.get('epochs_1st', 200) save_freq = config.get('save_freq', 2) log_interval = config.get('log_interval', 10) saving_epoch = config.get('save_freq', 2) data_params = config.get('data_params', None) sr = config['preprocess_params'].get('sr', 24000) train_path = data_params['train_data'] val_path = data_params['val_data'] root_path = data_params['root_path'] min_length = data_params['min_length'] OOD_data = data_params['OOD_data'] max_len = config.get('max_len', 200) # load data train_list, val_list = get_data_path_list(train_path, val_path) train_dataloader = build_dataloader(train_list, root_path, OOD_data=OOD_data, min_length=min_length, batch_size=batch_size, num_workers=2, dataset_config={}, device=device) val_dataloader = build_dataloader(val_list, root_path, OOD_data=OOD_data, min_length=min_length, batch_size=batch_size, validation=True, num_workers=0, device=device, dataset_config={}) with accelerator.main_process_first(): # load pretrained ASR model ASR_config = config.get('ASR_config', False) ASR_path = config.get('ASR_path', False) text_aligner = load_ASR_models(ASR_path, ASR_config) # load pretrained F0 model F0_path = config.get('F0_path', False) pitch_extractor = load_F0_models(F0_path) # load BERT model from Utils.PLBERT.util import load_plbert BERT_path = config.get('PLBERT_dir', False) plbert = load_plbert(BERT_path) scheduler_params = { "max_lr": float(config['optimizer_params'].get('lr', 1e-4)), "pct_start": float(config['optimizer_params'].get('pct_start', 0.0)), "epochs": epochs, "steps_per_epoch": len(train_dataloader), } model_params = recursive_munch(config['model_params']) multispeaker = model_params.multispeaker model = build_model(model_params, text_aligner, pitch_extractor, plbert) best_loss = float('inf') # best test loss loss_train_record = list([]) loss_test_record = list([]) loss_params = Munch(config['loss_params']) TMA_epoch = loss_params.TMA_epoch for k in model: model[k] = accelerator.prepare(model[k]) train_dataloader, val_dataloader = accelerator.prepare( train_dataloader, val_dataloader ) _ = [model[key].to(device) for key in model] # initialize optimizers after preparing models for compatibility with FSDP optimizer = build_optimizer({key: model[key].parameters() for key in model}, scheduler_params_dict={ key: scheduler_params.copy() for key in model}, lr=float(config['optimizer_params'].get('lr', 1e-4))) for k, v in optimizer.optimizers.items(): optimizer.optimizers[k] = accelerator.prepare(optimizer.optimizers[k]) optimizer.schedulers[k] = accelerator.prepare(optimizer.schedulers[k]) with accelerator.main_process_first(): if config.get('pretrained_model', '') != '': model, optimizer, start_epoch, iters = load_checkpoint(model, optimizer, config['pretrained_model'], load_only_params=config.get('load_only_params', True)) else: start_epoch = 0 iters = 0 # in case not distributed try: n_down = model.text_aligner.module.n_down except: n_down = model.text_aligner.n_down # wrapped losses for compatibility with mixed precision stft_loss = MultiResolutionSTFTLoss().to(device) gl = GeneratorLoss(model.mpd, model.msd).to(device) dl = DiscriminatorLoss(model.mpd, model.msd).to(device) wl = WavLMLoss(model_params.slm.model, model.wd, sr, model_params.slm.sr).to(device) for epoch in range(start_epoch, epochs): running_loss = 0 start_time = time.time() _ = [model[key].train() for key in model] for i, batch in enumerate(train_dataloader): waves = batch[0] batch = [b.to(device) for b in batch[1:]] texts, input_lengths, _, _, mels, mel_input_length, _ = batch with torch.no_grad(): mask = length_to_mask(mel_input_length // (2 ** n_down)).to('cuda') text_mask = length_to_mask(input_lengths).to(texts.device) ppgs, s2s_pred, s2s_attn = model.text_aligner(mels, mask, texts) s2s_attn = s2s_attn.transpose(-1, -2) s2s_attn = s2s_attn[..., 1:] s2s_attn = s2s_attn.transpose(-1, -2) with torch.no_grad(): attn_mask = (~mask).unsqueeze(-1).expand( mask.shape[0], mask.shape[1], text_mask.shape[-1]).float().transpose(-1, -2) attn_mask = attn_mask.float() * (~text_mask).unsqueeze(-1).expand( text_mask.shape[0], text_mask.shape[1], mask.shape[-1]).float() attn_mask = (attn_mask < 1) s2s_attn.masked_fill_(attn_mask, 0.0) with torch.no_grad(): mask_ST = mask_from_lens( s2s_attn, input_lengths, mel_input_length // (2 ** n_down)) s2s_attn_mono = maximum_path(s2s_attn, mask_ST) # encode t_en = model.text_encoder(texts, input_lengths, text_mask) # 50% of chance of using monotonic version if bool(random.getrandbits(1)): asr = (t_en @ s2s_attn) else: asr = (t_en @ s2s_attn_mono) # get clips mel_input_length_all = accelerator.gather( mel_input_length) # for balanced load mel_len = min( [int(mel_input_length_all.min().item() / 2 - 1), max_len // 2]) mel_len_st = int(mel_input_length.min().item() / 2 - 1) en = [] gt = [] wav = [] st = [] for bib in range(len(mel_input_length)): mel_length = int(mel_input_length[bib].item() / 2) random_start = np.random.randint(0, mel_length - mel_len) en.append(asr[bib, :, random_start:random_start+mel_len]) gt.append(mels[bib, :, (random_start * 2) :((random_start+mel_len) * 2)]) y = waves[bib][(random_start * 2) * 300:((random_start+mel_len) * 2) * 300] wav.append(torch.from_numpy(y).to(device)) # style reference (better to be different from the GT) random_start = np.random.randint(0, mel_length - mel_len_st) st.append(mels[bib, :, (random_start * 2) :((random_start+mel_len_st) * 2)]) en = torch.stack(en) gt = torch.stack(gt).detach() st = torch.stack(st).detach() wav = torch.stack(wav).float().detach() # clip too short to be used by the style encoder if gt.shape[-1] < 80: continue with torch.no_grad(): real_norm = log_norm(gt.unsqueeze(1)).squeeze(1).detach() F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1)) s = model.style_encoder(st.unsqueeze( 1) if multispeaker else gt.unsqueeze(1)) y_rec = model.decoder(en, F0_real, real_norm, s) # discriminator loss if epoch >= TMA_epoch: optimizer.zero_grad() d_loss = dl(wav.detach().unsqueeze( 1).float(), y_rec.detach()).mean() accelerator.backward(d_loss) optimizer.step('msd') optimizer.step('mpd') else: d_loss = 0 # generator loss optimizer.zero_grad() loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) if epoch >= TMA_epoch: # start TMA training loss_s2s = 0 for _s2s_pred, _text_input, _text_length in zip(s2s_pred, texts, input_lengths): loss_s2s += F.cross_entropy( _s2s_pred[:_text_length], _text_input[:_text_length]) loss_s2s /= texts.size(0) loss_mono = F.l1_loss(s2s_attn, s2s_attn_mono) * 10 loss_gen_all = gl(wav.detach().unsqueeze( 1).float(), y_rec).mean() loss_slm = wl(wav.detach(), y_rec).mean() g_loss = loss_params.lambda_mel * loss_mel + \ loss_params.lambda_mono * loss_mono + \ loss_params.lambda_s2s * loss_s2s + \ loss_params.lambda_gen * loss_gen_all + \ loss_params.lambda_slm * loss_slm else: loss_s2s = 0 loss_mono = 0 loss_gen_all = 0 loss_slm = 0 g_loss = loss_mel running_loss += accelerator.gather(loss_mel).mean().item() accelerator.backward(g_loss) optimizer.step('text_encoder') optimizer.step('style_encoder') optimizer.step('decoder') if epoch >= TMA_epoch: optimizer.step('text_aligner') optimizer.step('pitch_extractor') iters = iters + 1 if (i+1) % log_interval == 0 and accelerator.is_main_process: log_print('Epoch [%d/%d], Step [%d/%d], Mel Loss: %.5f, Gen Loss: %.5f, Disc Loss: %.5f, Mono Loss: %.5f, S2S Loss: %.5f, SLM Loss: %.5f' % (epoch+1, epochs, i+1, len(train_list)//batch_size, running_loss / log_interval, loss_gen_all, d_loss, loss_mono, loss_s2s, loss_slm), logger) writer.add_scalar('train/mel_loss', running_loss / log_interval, iters) writer.add_scalar('train/gen_loss', loss_gen_all, iters) writer.add_scalar('train/d_loss', d_loss, iters) writer.add_scalar('train/mono_loss', loss_mono, iters) writer.add_scalar('train/s2s_loss', loss_s2s, iters) writer.add_scalar('train/slm_loss', loss_slm, iters) running_loss = 0 print('Time elasped:', time.time()-start_time) loss_test = 0 _ = [model[key].eval() for key in model] with torch.no_grad(): iters_test = 0 for batch_idx, batch in enumerate(val_dataloader): optimizer.zero_grad() waves = batch[0] batch = [b.to(device) for b in batch[1:]] texts, input_lengths, _, _, mels, mel_input_length, _ = batch with torch.no_grad(): mask = length_to_mask( mel_input_length // (2 ** n_down)).to('cuda') ppgs, s2s_pred, s2s_attn = model.text_aligner( mels, mask, texts) s2s_attn = s2s_attn.transpose(-1, -2) s2s_attn = s2s_attn[..., 1:] s2s_attn = s2s_attn.transpose(-1, -2) text_mask = length_to_mask(input_lengths).to(texts.device) attn_mask = (~mask).unsqueeze(-1).expand( mask.shape[0], mask.shape[1], text_mask.shape[-1]).float().transpose(-1, -2) attn_mask = attn_mask.float() * (~text_mask).unsqueeze(-1).expand( text_mask.shape[0], text_mask.shape[1], mask.shape[-1]).float() attn_mask = (attn_mask < 1) s2s_attn.masked_fill_(attn_mask, 0.0) # encode t_en = model.text_encoder(texts, input_lengths, text_mask) asr = (t_en @ s2s_attn) # get clips mel_input_length_all = accelerator.gather( mel_input_length) # for balanced load mel_len = min( [int(mel_input_length.min().item() / 2 - 1), max_len // 2]) en = [] gt = [] wav = [] for bib in range(len(mel_input_length)): mel_length = int(mel_input_length[bib].item() / 2) random_start = np.random.randint(0, mel_length - mel_len) en.append(asr[bib, :, random_start:random_start+mel_len]) gt.append(mels[bib, :, (random_start * 2) :((random_start+mel_len) * 2)]) y = waves[bib][(random_start * 2) * 300:((random_start+mel_len) * 2) * 300] wav.append(torch.from_numpy(y).to(device)) wav = torch.stack(wav).float().detach() en = torch.stack(en) gt = torch.stack(gt).detach() if gt.shape[-1] < 80: continue with torch.no_grad(): real_norm = log_norm(gt.unsqueeze(1)).squeeze(1).detach() F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1)) s = model.style_encoder(st.unsqueeze( 1) if multispeaker else gt.unsqueeze(1)) y_rec = model.decoder(en, F0_real, real_norm, s) loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) loss_test += accelerator.gather(loss_mel).mean().item() iters_test += 1 if accelerator.is_main_process: print('Epochs:', epoch + 1) log_print('Validation loss: %.3f' % (loss_test / iters_test) + '\n\n\n\n', logger) print('\n\n\n') writer.add_scalar('eval/mel_loss', loss_test / iters_test, epoch + 1) attn_image = get_image(s2s_attn[0].cpu().numpy().squeeze()) writer.add_figure('eval/attn', attn_image, epoch) with torch.no_grad(): for bib in range(len(asr)): mel_length = int(mel_input_length[bib].item()) gt = mels[bib, :, :mel_length].unsqueeze(0) en = asr[bib, :, :mel_length // 2].unsqueeze(0) F0_real, _, _ = model.pitch_extractor(gt.unsqueeze(1)) F0_real = F0_real.unsqueeze(0) s = model.style_encoder(gt.unsqueeze(1)) real_norm = log_norm(gt.unsqueeze(1)).squeeze(1) y_rec = model.decoder(en, F0_real, real_norm, s) writer.add_audio( 'eval/y' + str(bib), y_rec.cpu().numpy().squeeze(), epoch, sample_rate=sr) if epoch == 0: writer.add_audio( 'gt/y' + str(bib), waves[bib].squeeze(), epoch, sample_rate=sr) if bib >= 6: break if epoch % saving_epoch == 0: if (loss_test / iters_test) < best_loss: best_loss = loss_test / iters_test print('Saving..') state = { 'net': {key: model[key].state_dict() for key in model}, 'optimizer': optimizer.state_dict(), 'iters': iters, 'val_loss': loss_test / iters_test, 'epoch': epoch, } save_path = osp.join(log_dir, 'epoch_1st_%05d.pth' % epoch) torch.save(state, save_path) if accelerator.is_main_process: print('Saving..') state = { 'net': {key: model[key].state_dict() for key in model}, 'optimizer': optimizer.state_dict(), 'iters': iters, 'val_loss': loss_test / iters_test, 'epoch': epoch, } save_path = osp.join(log_dir, config.get( 'first_stage_path', 'first_stage.pth')) torch.save(state, save_path) if __name__ == "__main__": main()