import os, sys # sys.path.append('/home/shaunxliu/projects/nnsp') import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.ticker import MaxNLocator import torch from torch.utils.data import DataLoader import numpy as np from .solver import BaseSolver from utils.data_load import OneshotVcDataset, MultiSpkVcCollate # from src.rnn_ppg2mel import BiRnnPpg2MelModel # from src.mel_decoder_mol_encAddlf0 import MelDecoderMOL from .loss import MaskedMSELoss from .optim import Optimizer from utils.util import human_format from models.ppg2mel import MelDecoderMOLv2 class Solver(BaseSolver): """Customized Solver.""" def __init__(self, config, paras, mode): super().__init__(config, paras, mode) self.num_att_plots = 5 self.att_ws_dir = f"{self.logdir}/att_ws" os.makedirs(self.att_ws_dir, exist_ok=True) self.best_loss = np.inf def fetch_data(self, data): """Move data to device""" data = [i.to(self.device) for i in data] return data def load_data(self): """ Load data for training/validation/plotting.""" train_dataset = OneshotVcDataset( meta_file=self.config.data.train_fid_list, vctk_ppg_dir=self.config.data.vctk_ppg_dir, libri_ppg_dir=self.config.data.libri_ppg_dir, vctk_f0_dir=self.config.data.vctk_f0_dir, libri_f0_dir=self.config.data.libri_f0_dir, vctk_wav_dir=self.config.data.vctk_wav_dir, libri_wav_dir=self.config.data.libri_wav_dir, vctk_spk_dvec_dir=self.config.data.vctk_spk_dvec_dir, libri_spk_dvec_dir=self.config.data.libri_spk_dvec_dir, ppg_file_ext=self.config.data.ppg_file_ext, min_max_norm_mel=self.config.data.min_max_norm_mel, mel_min=self.config.data.mel_min, mel_max=self.config.data.mel_max, ) dev_dataset = OneshotVcDataset( meta_file=self.config.data.dev_fid_list, vctk_ppg_dir=self.config.data.vctk_ppg_dir, libri_ppg_dir=self.config.data.libri_ppg_dir, vctk_f0_dir=self.config.data.vctk_f0_dir, libri_f0_dir=self.config.data.libri_f0_dir, vctk_wav_dir=self.config.data.vctk_wav_dir, libri_wav_dir=self.config.data.libri_wav_dir, vctk_spk_dvec_dir=self.config.data.vctk_spk_dvec_dir, libri_spk_dvec_dir=self.config.data.libri_spk_dvec_dir, ppg_file_ext=self.config.data.ppg_file_ext, min_max_norm_mel=self.config.data.min_max_norm_mel, mel_min=self.config.data.mel_min, mel_max=self.config.data.mel_max, ) self.train_dataloader = DataLoader( train_dataset, num_workers=self.paras.njobs, shuffle=True, batch_size=self.config.hparas.batch_size, pin_memory=False, drop_last=True, collate_fn=MultiSpkVcCollate(self.config.model.frames_per_step, use_spk_dvec=True), ) self.dev_dataloader = DataLoader( dev_dataset, num_workers=self.paras.njobs, shuffle=False, batch_size=self.config.hparas.batch_size, pin_memory=False, drop_last=False, collate_fn=MultiSpkVcCollate(self.config.model.frames_per_step, use_spk_dvec=True), ) self.plot_dataloader = DataLoader( dev_dataset, num_workers=self.paras.njobs, shuffle=False, batch_size=1, pin_memory=False, drop_last=False, collate_fn=MultiSpkVcCollate(self.config.model.frames_per_step, use_spk_dvec=True, give_uttids=True), ) msg = "Have prepared training set and dev set." self.verbose(msg) def load_pretrained_params(self): print("Load pretrained model from: ", self.config.data.pretrain_model_file) ignore_layer_prefixes = ["speaker_embedding_table"] pretrain_model_file = self.config.data.pretrain_model_file pretrain_ckpt = torch.load( pretrain_model_file, map_location=self.device )["model"] model_dict = self.model.state_dict() print(self.model) # 1. filter out unnecessrary keys for prefix in ignore_layer_prefixes: pretrain_ckpt = {k : v for k, v in pretrain_ckpt.items() if not k.startswith(prefix) } # 2. overwrite entries in the existing state dict model_dict.update(pretrain_ckpt) # 3. load the new state dict self.model.load_state_dict(model_dict) def set_model(self): """Setup model and optimizer""" # Model print("[INFO] Model name: ", self.config["model_name"]) self.model = MelDecoderMOLv2( **self.config["model"] ).to(self.device) # self.load_pretrained_params() # model_params = [{'params': self.model.spk_embedding.weight}] model_params = [{'params': self.model.parameters()}] # Loss criterion self.loss_criterion = MaskedMSELoss(self.config.model.frames_per_step) # Optimizer self.optimizer = Optimizer(model_params, **self.config["hparas"]) self.verbose(self.optimizer.create_msg()) # Automatically load pre-trained model if self.paras.load is given self.load_ckpt() def exec(self): self.verbose("Total training steps {}.".format( human_format(self.max_step))) mel_loss = None n_epochs = 0 # Set as current time self.timer.set() while self.step < self.max_step: for data in self.train_dataloader: # Pre-step: updata lr_rate and do zero_grad lr_rate = self.optimizer.pre_step(self.step) total_loss = 0 # data to device ppgs, lf0_uvs, mels, in_lengths, \ out_lengths, spk_ids, stop_tokens = self.fetch_data(data) self.timer.cnt("rd") mel_outputs, mel_outputs_postnet, predicted_stop = self.model( ppgs, in_lengths, mels, out_lengths, lf0_uvs, spk_ids ) mel_loss, stop_loss = self.loss_criterion( mel_outputs, mel_outputs_postnet, mels, out_lengths, stop_tokens, predicted_stop ) loss = mel_loss + stop_loss self.timer.cnt("fw") # Back-prop grad_norm = self.backward(loss) self.step += 1 # Logger if (self.step == 1) or (self.step % self.PROGRESS_STEP == 0): self.progress("Tr|loss:{:.4f},mel-loss:{:.4f},stop-loss:{:.4f}|Grad.Norm-{:.2f}|{}" .format(loss.cpu().item(), mel_loss.cpu().item(), stop_loss.cpu().item(), grad_norm, self.timer.show())) self.write_log('loss', {'tr/loss': loss, 'tr/mel-loss': mel_loss, 'tr/stop-loss': stop_loss}) # Validation if (self.step == 1) or (self.step % self.valid_step == 0): self.validate() # End of step # https://github.com/pytorch/pytorch/issues/13246#issuecomment-529185354 torch.cuda.empty_cache() self.timer.set() if self.step > self.max_step: break n_epochs += 1 self.log.close() def validate(self): self.model.eval() dev_loss, dev_mel_loss, dev_stop_loss = 0.0, 0.0, 0.0 for i, data in enumerate(self.dev_dataloader): self.progress('Valid step - {}/{}'.format(i+1, len(self.dev_dataloader))) # Fetch data ppgs, lf0_uvs, mels, in_lengths, \ out_lengths, spk_ids, stop_tokens = self.fetch_data(data) with torch.no_grad(): mel_outputs, mel_outputs_postnet, predicted_stop = self.model( ppgs, in_lengths, mels, out_lengths, lf0_uvs, spk_ids ) mel_loss, stop_loss = self.loss_criterion( mel_outputs, mel_outputs_postnet, mels, out_lengths, stop_tokens, predicted_stop ) loss = mel_loss + stop_loss dev_loss += loss.cpu().item() dev_mel_loss += mel_loss.cpu().item() dev_stop_loss += stop_loss.cpu().item() dev_loss = dev_loss / (i + 1) dev_mel_loss = dev_mel_loss / (i + 1) dev_stop_loss = dev_stop_loss / (i + 1) self.save_checkpoint(f'step_{self.step}.pth', 'loss', dev_loss, show_msg=False) if dev_loss < self.best_loss: self.best_loss = dev_loss self.save_checkpoint(f'best_loss_step_{self.step}.pth', 'loss', dev_loss) self.write_log('loss', {'dv/loss': dev_loss, 'dv/mel-loss': dev_mel_loss, 'dv/stop-loss': dev_stop_loss}) # plot attention for i, data in enumerate(self.plot_dataloader): if i == self.num_att_plots: break # Fetch data ppgs, lf0_uvs, mels, in_lengths, \ out_lengths, spk_ids, stop_tokens = self.fetch_data(data[:-1]) fid = data[-1][0] with torch.no_grad(): _, _, _, att_ws = self.model( ppgs, in_lengths, mels, out_lengths, lf0_uvs, spk_ids, output_att_ws=True ) att_ws = att_ws.squeeze(0).cpu().numpy() att_ws = att_ws[None] w, h = plt.figaspect(1.0 / len(att_ws)) fig = plt.Figure(figsize=(w * 1.3, h * 1.3)) axes = fig.subplots(1, len(att_ws)) if len(att_ws) == 1: axes = [axes] for ax, aw in zip(axes, att_ws): ax.imshow(aw.astype(np.float32), aspect="auto") ax.set_title(f"{fid}") ax.set_xlabel("Input") ax.set_ylabel("Output") ax.xaxis.set_major_locator(MaxNLocator(integer=True)) ax.yaxis.set_major_locator(MaxNLocator(integer=True)) fig_name = f"{self.att_ws_dir}/{fid}_step{self.step}.png" fig.savefig(fig_name) # Resume training self.model.train()