import filecmp import os import traceback import numpy as np import pandas as pd import torch import torch.distributed as dist import torch.nn.functional as F import torch.optim import torch.utils.data import yaml from tqdm import tqdm import utils from tasks.tts.dataset_utils import BaseSpeechDataset from tasks.tts.tts_utils import parse_mel_losses, parse_dataset_configs, load_data_preprocessor, load_data_binarizer from tasks.tts.vocoder_infer.base_vocoder import BaseVocoder, get_vocoder_cls from utils.audio.align import mel2token_to_dur from utils.audio.io import save_wav from utils.audio.pitch_extractors import extract_pitch_simple from utils.commons.base_task import BaseTask from utils.commons.ckpt_utils import load_ckpt from utils.commons.dataset_utils import data_loader, BaseConcatDataset from utils.commons.hparams import hparams from utils.commons.multiprocess_utils import MultiprocessManager from utils.commons.tensor_utils import tensors_to_scalars from utils.metrics.ssim import ssim from utils.nn.model_utils import print_arch from utils.nn.schedulers import RSQRTSchedule, NoneSchedule, WarmupSchedule from utils.nn.seq_utils import weights_nonzero_speech from utils.plot.plot import spec_to_figure from utils.text.text_encoder import build_token_encoder import matplotlib.pyplot as plt class SpeechBaseTask(BaseTask): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.dataset_cls = BaseSpeechDataset self.vocoder = None data_dir = hparams['binary_data_dir'] if not hparams['use_word_input']: self.token_encoder = build_token_encoder(f'{data_dir}/phone_set.json') else: self.token_encoder = build_token_encoder(f'{data_dir}/word_set.json') self.padding_idx = self.token_encoder.pad() self.eos_idx = self.token_encoder.eos() self.seg_idx = self.token_encoder.seg() self.saving_result_pool = None self.saving_results_futures = None self.mel_losses = parse_mel_losses() self.max_tokens, self.max_sentences, \ self.max_valid_tokens, self.max_valid_sentences = parse_dataset_configs() ########################## # datasets ########################## @data_loader def train_dataloader(self): if hparams['train_sets'] != '': train_sets = hparams['train_sets'].split("|") # check if all train_sets have the same spk map and dictionary binary_data_dir = hparams['binary_data_dir'] file_to_cmp = ['phone_set.json'] if os.path.exists(f'{binary_data_dir}/word_set.json'): file_to_cmp.append('word_set.json') if hparams['use_spk_id']: file_to_cmp.append('spk_map.json') for f in file_to_cmp: for ds_name in train_sets: base_file = os.path.join(binary_data_dir, f) ds_file = os.path.join(ds_name, f) assert filecmp.cmp(base_file, ds_file), \ f'{f} in {ds_name} is not same with that in {binary_data_dir}.' train_dataset = BaseConcatDataset([ self.dataset_cls(prefix='train', shuffle=True, data_dir=ds_name) for ds_name in train_sets]) else: train_dataset = self.dataset_cls(prefix=hparams['train_set_name'], shuffle=True) return self.build_dataloader(train_dataset, True, self.max_tokens, self.max_sentences, endless=hparams['endless_ds']) @data_loader def val_dataloader(self): valid_dataset = self.dataset_cls(prefix=hparams['valid_set_name'], shuffle=False) return self.build_dataloader(valid_dataset, False, self.max_valid_tokens, self.max_valid_sentences, batch_by_size=False) @data_loader def test_dataloader(self): test_dataset = self.dataset_cls(prefix=hparams['test_set_name'], shuffle=False) self.test_dl = self.build_dataloader( test_dataset, False, self.max_valid_tokens, self.max_valid_sentences, batch_by_size=False) return self.test_dl def build_dataloader(self, dataset, shuffle, max_tokens=None, max_sentences=None, required_batch_size_multiple=-1, endless=False, batch_by_size=True): devices_cnt = torch.cuda.device_count() if devices_cnt == 0: devices_cnt = 1 if required_batch_size_multiple == -1: required_batch_size_multiple = devices_cnt def shuffle_batches(batches): np.random.shuffle(batches) return batches if max_tokens is not None: max_tokens *= devices_cnt if max_sentences is not None: max_sentences *= devices_cnt indices = dataset.ordered_indices() if batch_by_size: batch_sampler = utils.commons.dataset_utils.batch_by_size( indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences, required_batch_size_multiple=required_batch_size_multiple, ) else: batch_sampler = [] for i in range(0, len(indices), max_sentences): batch_sampler.append(indices[i:i + max_sentences]) if shuffle: batches = shuffle_batches(list(batch_sampler)) if endless: batches = [b for _ in range(1000) for b in shuffle_batches(list(batch_sampler))] else: batches = batch_sampler if endless: batches = [b for _ in range(1000) for b in batches] num_workers = dataset.num_workers if self.trainer.use_ddp: num_replicas = dist.get_world_size() rank = dist.get_rank() batches = [x[rank::num_replicas] for x in batches if len(x) % num_replicas == 0] return torch.utils.data.DataLoader(dataset, collate_fn=dataset.collater, batch_sampler=batches, num_workers=num_workers, pin_memory=False) ########################## # scheduler and optimizer ########################## def build_model(self): self.build_tts_model() if hparams['load_ckpt'] != '': load_ckpt(self.model, hparams['load_ckpt']) print_arch(self.model) return self.model def build_tts_model(self): raise NotImplementedError def build_scheduler(self, optimizer): if hparams['scheduler'] == 'rsqrt': return RSQRTSchedule(optimizer, hparams['lr'], hparams['warmup_updates'], hparams['hidden_size']) elif hparams['scheduler'] == 'warmup': return WarmupSchedule(optimizer, hparams['lr'], hparams['warmup_updates']) elif hparams['scheduler'] == 'step_lr': return torch.optim.lr_scheduler.StepLR( optimizer=optimizer, step_size=500, gamma=0.998) else: return NoneSchedule(optimizer, hparams['lr']) def build_optimizer(self, model): self.optimizer = optimizer = torch.optim.AdamW( model.parameters(), lr=hparams['lr'], betas=(hparams['optimizer_adam_beta1'], hparams['optimizer_adam_beta2']), weight_decay=hparams['weight_decay']) return optimizer ########################## # training and validation ########################## def _training_step(self, sample, batch_idx, _): loss_output, _ = self.run_model(sample) total_loss = sum([v for v in loss_output.values() if isinstance(v, torch.Tensor) and v.requires_grad]) loss_output['batch_size'] = sample['txt_tokens'].size()[0] return total_loss, loss_output def run_model(self, sample, infer=False): """ :param sample: a batch of data :param infer: bool, run in infer mode :return: if not infer: return losses, model_out if infer: return model_out """ raise NotImplementedError def validation_start(self): self.vocoder = get_vocoder_cls(hparams['vocoder'])() def validation_step(self, sample, batch_idx): outputs = {} outputs['losses'] = {} outputs['losses'], model_out = self.run_model(sample) outputs['total_loss'] = sum(outputs['losses'].values()) outputs['nsamples'] = sample['nsamples'] outputs = tensors_to_scalars(outputs) if self.global_step % hparams['valid_infer_interval'] == 0 \ and batch_idx < hparams['num_valid_plots']: self.save_valid_result(sample, batch_idx, model_out) return outputs def validation_end(self, outputs): self.vocoder = None return super(SpeechBaseTask, self).validation_end(outputs) def save_valid_result(self, sample, batch_idx, model_out): raise NotImplementedError ########################## # losses ########################## def add_mel_loss(self, mel_out, target, losses, postfix=''): for loss_name, lambd in self.mel_losses.items(): losses[f'{loss_name}{postfix}'] = getattr(self, f'{loss_name}_loss')(mel_out, target) * lambd def l1_loss(self, decoder_output, target): # decoder_output : B x T x n_mel # target : B x T x n_mel l1_loss = F.l1_loss(decoder_output, target, reduction='none') weights = weights_nonzero_speech(target) l1_loss = (l1_loss * weights).sum() / weights.sum() return l1_loss def mse_loss(self, decoder_output, target): # decoder_output : B x T x n_mel # target : B x T x n_mel assert decoder_output.shape == target.shape mse_loss = F.mse_loss(decoder_output, target, reduction='none') weights = weights_nonzero_speech(target) mse_loss = (mse_loss * weights).sum() / weights.sum() return mse_loss def ssim_loss(self, decoder_output, target, bias=6.0): # decoder_output : B x T x n_mel # target : B x T x n_mel assert decoder_output.shape == target.shape weights = weights_nonzero_speech(target) decoder_output = decoder_output[:, None] + bias target = target[:, None] + bias ssim_loss = 1 - ssim(decoder_output, target, size_average=False) ssim_loss = (ssim_loss * weights).sum() / weights.sum() return ssim_loss def plot_mel(self, batch_idx, spec_out, spec_gt=None, name=None, title='', f0s=None, dur_info=None): vmin = hparams['mel_vmin'] vmax = hparams['mel_vmax'] if len(spec_out.shape) == 3: spec_out = spec_out[0] if isinstance(spec_out, torch.Tensor): spec_out = spec_out.cpu().numpy() if spec_gt is not None: if len(spec_gt.shape) == 3: spec_gt = spec_gt[0] if isinstance(spec_gt, torch.Tensor): spec_gt = spec_gt.cpu().numpy() max_len = max(len(spec_gt), len(spec_out)) if max_len - len(spec_gt) > 0: spec_gt = np.pad(spec_gt, [[0, max_len - len(spec_gt)], [0, 0]], mode='constant', constant_values=vmin) if max_len - len(spec_out) > 0: spec_out = np.pad(spec_out, [[0, max_len - len(spec_out)], [0, 0]], mode='constant', constant_values=vmin) spec_out = np.concatenate([spec_out, spec_gt], -1) name = f'mel_val_{batch_idx}' if name is None else name self.logger.add_figure(name, spec_to_figure( spec_out, vmin, vmax, title=title, f0s=f0s, dur_info=dur_info), self.global_step) ########################## # testing ########################## def test_start(self): self.saving_result_pool = MultiprocessManager(int(os.getenv('N_PROC', os.cpu_count()))) self.saving_results_futures = [] self.gen_dir = os.path.join( hparams['work_dir'], f'generated_{self.trainer.global_step}_{hparams["gen_dir_name"]}') self.vocoder: BaseVocoder = get_vocoder_cls(hparams['vocoder'])() os.makedirs(self.gen_dir, exist_ok=True) os.makedirs(f'{self.gen_dir}/wavs', exist_ok=True) os.makedirs(f'{self.gen_dir}/plot', exist_ok=True) if hparams.get('save_mel_npy', False): os.makedirs(f'{self.gen_dir}/mel_npy', exist_ok=True) def test_step(self, sample, batch_idx): """ :param sample: :param batch_idx: :return: """ assert sample['txt_tokens'].shape[0] == 1, 'only support batch_size=1 in inference' outputs = self.run_model(sample, infer=True) text = sample['text'][0] item_name = sample['item_name'][0] tokens = sample['txt_tokens'][0].cpu().numpy() mel_gt = sample['mels'][0].cpu().numpy() mel_pred = outputs['mel_out'][0].cpu().numpy() str_phs = self.token_encoder.decode(tokens, strip_padding=True) base_fn = f'[{self.results_id:06d}][{item_name.replace("%", "_")}][%s]' if text is not None: base_fn += text.replace(":", "$3A")[:80] base_fn = base_fn.replace(' ', '_') gen_dir = self.gen_dir wav_pred = self.vocoder.spec2wav(mel_pred) self.saving_result_pool.add_job(self.save_result, args=[ wav_pred, mel_pred, base_fn % 'P', gen_dir, str_phs]) if hparams['save_gt']: wav_gt = self.vocoder.spec2wav(mel_gt) self.saving_result_pool.add_job(self.save_result, args=[ wav_gt, mel_gt, base_fn % 'G', gen_dir, str_phs]) print(f"Pred_shape: {mel_pred.shape}, gt_shape: {mel_gt.shape}") return { 'item_name': item_name, 'text': text, 'ph_tokens': self.token_encoder.decode(tokens.tolist()), 'wav_fn_pred': base_fn % 'P', 'wav_fn_gt': base_fn % 'G', } @staticmethod def save_result(wav_out, mel, base_fn, gen_dir, str_phs=None, mel2ph=None, alignment=None): save_wav(wav_out, f'{gen_dir}/wavs/{base_fn}.wav', hparams['audio_sample_rate'], norm=hparams['out_wav_norm']) fig = plt.figure(figsize=(14, 10)) spec_vmin = hparams['mel_vmin'] spec_vmax = hparams['mel_vmax'] heatmap = plt.pcolor(mel.T, vmin=spec_vmin, vmax=spec_vmax) fig.colorbar(heatmap) try: f0 = extract_pitch_simple(wav_out) f0 = f0 / 10 * (f0 > 0) plt.plot(f0, c='white', linewidth=1, alpha=0.6) if mel2ph is not None and str_phs is not None: decoded_txt = str_phs.split(" ") dur = mel2token_to_dur(torch.LongTensor(mel2ph)[None, :], len(decoded_txt))[0].numpy() dur = [0] + list(np.cumsum(dur)) for i in range(len(dur) - 1): shift = (i % 20) + 1 plt.text(dur[i], shift, decoded_txt[i]) plt.hlines(shift, dur[i], dur[i + 1], colors='b' if decoded_txt[i] != '|' else 'black') plt.vlines(dur[i], 0, 5, colors='b' if decoded_txt[i] != '|' else 'black', alpha=1, linewidth=1) plt.tight_layout() plt.savefig(f'{gen_dir}/plot/{base_fn}.png', format='png') plt.close(fig) if hparams.get('save_mel_npy', False): np.save(f'{gen_dir}/mel_npy/{base_fn}', mel) if alignment is not None: fig, ax = plt.subplots(figsize=(12, 16)) im = ax.imshow(alignment, aspect='auto', origin='lower', interpolation='none') decoded_txt = str_phs.split(" ") ax.set_yticks(np.arange(len(decoded_txt))) ax.set_yticklabels(list(decoded_txt), fontsize=6) fig.colorbar(im, ax=ax) fig.savefig(f'{gen_dir}/attn_plot/{base_fn}_attn.png', format='png') plt.close(fig) except Exception: traceback.print_exc() return None def test_end(self, outputs): pd.DataFrame(outputs).to_csv(f'{self.gen_dir}/meta.csv') for _1, _2 in tqdm(self.saving_result_pool.get_results(), total=len(self.saving_result_pool)): pass return {}