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 import random import yaml from munch import Munch import numpy as np import torch from torch import nn import torch.nn.functional as F import torchaudio import librosa from models import * from meldataset import build_dataloader from utils import * from losses import * from optimizers import build_optimizer import time from accelerate import Accelerator from accelerate.utils import LoggerType from accelerate import DistributedDataParallelKwargs from torch.utils.tensorboard import SummaryWriter import logging from accelerate.logging import get_logger 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)) save_iter = 10500 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], mixed_precision='bf16') 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) if (i+1)%save_iter == 0 and accelerator.is_main_process: print(f'Saving on step {epoch*len(train_dataloader)+i}...') state = { 'net': {key: model[key].state_dict() for key in model}, 'optimizer': optimizer.state_dict(), 'iters': iters, 'epoch': epoch, } save_path = osp.join(log_dir, f'2nd_phase_{epoch*len(train_dataloader)+i}.pth') torch.save(state, save_path) 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('cuda')) wav = torch.stack(wav).float().detach() en = torch.stack(en) gt = torch.stack(gt).detach() F0_real, _, F0 = model.pitch_extractor(gt.unsqueeze(1)) 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) 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 >= 15: 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()