# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import itertools from typing import Iterable, Optional import librosa import torch from kaldialign import edit_distance from lightning.pytorch import Trainer from lightning.pytorch.loggers import TensorBoardLogger from lightning.pytorch.utilities.combined_loader import CombinedLoader from omegaconf import DictConfig, OmegaConf from nemo.collections.asr.losses.angularloss import AngularSoftmaxLoss from nemo.collections.tts.data.dataset import TTSDataset from nemo.collections.tts.modules.ssl_tts import GreedyCTCDecoder from nemo.collections.tts.torch.tts_tokenizers import BaseTokenizer, EnglishCharsTokenizer from nemo.core.classes import ModelPT from nemo.core.classes.common import PretrainedModelInfo, safe_instantiate from nemo.core.optim.lr_scheduler import WarmupPolicy from nemo.utils import logging from nemo.utils.decorators import experimental @experimental class SSLDisentangler(ModelPT): """ SSLDisentangler is a Conformer based model for extracting disentangled content and speaker embeddings from an audio waveform. This model uses a pre-trained Conformer SSL model. To extract the linguistic content and speaker representations using a pre-trained Conformer, two randomly initialized downstream heads are added and the entire setup is finetuned in multi-task manner for speech recognition and speaker verification. These representations can be used by FastPitchModel_SSL for voice conversion by swapping the speaker embedding of a given source utterance, with the speaker embedding of a target speaker. """ def __init__(self, cfg: DictConfig, trainer: Trainer = None): super().__init__(cfg=cfg, trainer=trainer) self.preprocessor_disentangler = SSLDisentangler.from_config_dict(self._cfg.preprocessor) self.encoder = SSLDisentangler.from_config_dict(self._cfg.encoder) self._text_tokenizer = EnglishCharsTokenizer(add_blank_at="last") self._tb_logger = None self.downstream_nets = torch.nn.ModuleDict() for task in self._cfg.downstream_heads.task_names: if task == 'speaker_verification': # setting up downstream heads and loss functions for speaker verification task in_dim = self._cfg.encoder.d_model out_dim = self._cfg.downstream_heads.speaker_embed_size num_speakers = self._cfg.downstream_heads.num_speakers self.downstream_nets[task] = torch.nn.Linear(in_dim, out_dim) self.sv_linear = torch.nn.Linear(out_dim, num_speakers) self.sv_loss = AngularSoftmaxLoss(scale=30, margin=0.4) elif task == 'content': # setting up downstream heads and loss functions for text/content recognition task in_dim = self._cfg.encoder.d_model out_dim = self._cfg.downstream_heads.content_embed_size num_chars = len(self._text_tokenizer.tokens) # list of english tokens self.downstream_nets[task] = torch.nn.Linear(in_dim, out_dim) self.content_linear = torch.nn.Linear(out_dim, num_chars) self.ctc_loss = torch.nn.CTCLoss(blank=self._text_tokenizer.blank, zero_infinity=True) self.pitch_augment = self._cfg.get('pitch_augment', False) self.augment_ctc = self._cfg.get('augment_ctc', False) self.aug_loss_type = self._cfg.get('aug_loss_type', 'mse') self.stop_gradient = self._cfg.get('stop_gradient', False) assert ( self.stop_gradient and self.augment_ctc ) == False, "stop_gradient and augment_ctc cannot be true at the same time" self.mse_loss = torch.nn.MSELoss() self.ctc_decoder = GreedyCTCDecoder(self._text_tokenizer.tokens, self._text_tokenizer.blank) else: raise ValueError(f"{task} is not a valid task. Task must be speaker_verification or content.") self.automatic_optimization = False stft_cfg = self._cfg.preprocessor librosa_mel_filter = librosa.filters.mel( sr=stft_cfg.sample_rate, n_fft=stft_cfg.n_fft, n_mels=stft_cfg.features, fmin=0, fmax=8000 ) fb = torch.tensor( librosa_mel_filter, dtype=torch.float, ).unsqueeze(0) self.register_buffer("fb", fb) @classmethod def list_available_models(cls) -> Optional[PretrainedModelInfo]: """ This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud. Returns: List of available pre-trained models. """ results = [] model = PretrainedModelInfo( pretrained_model_name="ssl_en_conformer_large", description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:ssl_en_conformer_large", location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/ssl_en_conformer_large/versions/1.10.1/files/ssl_en_conformer_large.nemo", ) results.append(model) model = PretrainedModelInfo( pretrained_model_name="ssl_en_conformer_xlarge", description="For details about this model, please visit https://ngc.nvidia.com/catalog/models/nvidia:nemo:ssl_en_conformer_xlarge", location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/ssl_en_conformer_xlarge/versions/1.10.0/files/ssl_en_conformer_xlarge.nemo", ) results.append(model) return results @property def tb_logger(self): if self._tb_logger is None: if self.logger is None and self.logger.experiment is None: return None tb_logger = self.logger.experiment if isinstance(self.logger, Iterable): for logger in self.logger: if isinstance(logger, TensorBoardLogger): tb_logger = logger.experiment break self._tb_logger = tb_logger return self._tb_logger def __setup_dataloader_from_config(self, data_config): if hasattr(self, '_text_tokenizer') and isinstance(self._text_tokenizer, BaseTokenizer): _text_tokenizer = self._text_tokenizer else: if hasattr(self, '_text_tokenizer') and not isinstance(self._text_tokenizer, BaseTokenizer): logging.warning("test_tokenizer is set but not a BaseTokenizer. Will be set to EnglishCharsTokenizer") _text_tokenizer = self._text_tokenizer = EnglishCharsTokenizer(add_blank_at="last") for task in self._cfg.downstream_heads.task_names: if task == 'speaker_verification': sv_dataset = TTSDataset( manifest_filepath=data_config['manifest_speaker_verification_fp'], sample_rate=self._cfg.sample_rate, text_tokenizer=_text_tokenizer, segment_max_duration=data_config['segment_max_duration'], sup_data_types=['speaker_id'], sup_data_path=data_config['sup_data_path'], pad_multiple=data_config.get('pad_multiple', 1), ) sv_loader = torch.utils.data.DataLoader( sv_dataset, batch_size=data_config['batch_size_sv'], collate_fn=sv_dataset.general_collate_fn, shuffle=data_config['shuffle'], num_workers=data_config.get('num_workers_sv', 0), pin_memory=data_config.get('pin_memory', False), ) elif task == 'content': content_dataset = TTSDataset( manifest_filepath=data_config['manifest_content_fp'], sample_rate=self._cfg.sample_rate, text_tokenizer=_text_tokenizer, min_duration=data_config['min_duration_content'], max_duration=data_config['max_duration_content'], pitch_augment=data_config.get('pitch_augment', False), cache_pitch_augment=data_config.get('cache_pitch_augment', True), sup_data_path=data_config['sup_data_path'], pad_multiple=data_config.get('pad_multiple', 1), ) content_loader = torch.utils.data.DataLoader( content_dataset, batch_size=data_config['batch_size_content'], collate_fn=content_dataset.general_collate_fn, shuffle=data_config['shuffle'], num_workers=data_config.get('num_workers_content', 0), pin_memory=data_config.get('pin_memory', False), ) else: raise ValueError(f"{task} is not a valid task. Task must be speaker_verification or content.") loaders = {"sv": sv_loader, "content": content_loader} return loaders def setup_training_data(self, cfg): self._train_dl = self.__setup_dataloader_from_config(self._cfg.train_ds) def setup_validation_data(self, cfg): self._validation_dl = CombinedLoader(self.__setup_dataloader_from_config(self._cfg.validation_ds)) def configure_optimizers(self): optim_backbone_config = self._cfg.optim_backbone.copy() optim_downstream_config = self._cfg.optim_downstream.copy() OmegaConf.set_struct(optim_backbone_config, False) sched_backbone_config = optim_backbone_config.pop("sched", None) OmegaConf.set_struct(optim_backbone_config, True) OmegaConf.set_struct(optim_downstream_config, False) sched_downstream_config = optim_downstream_config.pop("sched", None) OmegaConf.set_struct(optim_downstream_config, True) optim_backbone = safe_instantiate( optim_backbone_config, params=self.encoder.parameters(), ) optim_downstream = safe_instantiate( optim_downstream_config, params=itertools.chain( self.downstream_nets.parameters(), self.sv_linear.parameters(), self.content_linear.parameters(), self.sv_loss.parameters(), ), ) if sched_backbone_config is not None and sched_downstream_config is not None: scheduler_backbone = WarmupPolicy( optimizer=optim_backbone, max_steps=None, min_lr=sched_backbone_config.min_lr, warmup_steps=sched_backbone_config.warmup_steps, ) # Use warmup to delay start sch1_dict = { 'scheduler': scheduler_backbone, 'interval': 'step', } scheduler_downstream = WarmupPolicy( optimizer=optim_downstream, max_steps=None, min_lr=sched_downstream_config.min_lr, warmup_steps=sched_downstream_config.warmup_steps, ) sch2_dict = { 'scheduler': scheduler_downstream, 'interval': 'step', } return [optim_backbone, optim_downstream], [sch1_dict, sch2_dict] else: return [optim_backbone, optim_downstream] def forward(self, input_signal=None, input_signal_length=None, normalize_content=True): processed_signal, processed_signal_length = self.preprocessor_disentangler( input_signal=input_signal, length=input_signal_length, ) encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length) # b,c,t for task in self._cfg.downstream_heads.task_names: if task == "speaker_verification": speaker_embedding = self.downstream_nets['speaker_verification'](encoded[:, :, 0]) l2_norm = torch.norm(speaker_embedding, p=2, dim=-1, keepdim=True) speaker_embedding_normalized = speaker_embedding / l2_norm speaker_logits = self.sv_linear(speaker_embedding_normalized) elif task == "content": encoded_btc = encoded.permute(0, 2, 1) content_embedding = self.downstream_nets['content'](encoded_btc) if normalize_content: l2_norm_content = torch.norm(content_embedding, p=2, dim=-1, keepdim=True) content_embedding = content_embedding / l2_norm_content content_logits = self.content_linear(content_embedding) content_log_probs = content_logits.log_softmax(dim=2) content_log_probs = content_log_probs.permute(1, 0, 2) # t,b,c for ctc else: raise ValueError(f"{task} is not a valid task. Task must be speaker_verification or content.") return ( speaker_logits, speaker_embedding_normalized, content_embedding, content_log_probs, encoded_len, ) def forward_for_export(self, input_signal=None, input_signal_length=None, normalize_content=True): # Same as forward right now. Earlier version of encoder had a different forward for export. # This function is still kept for compatibility with older evaluation/inference scripts. return self.forward( input_signal=input_signal, input_signal_length=input_signal_length, normalize_content=normalize_content, ) def training_step(self, batch, batch_idx): loss = 0.0 optim_backbone, optim_downstream = self.optimizers() schedulers = self.lr_schedulers() for key in batch.keys(): if key == 'sv': signal = batch[key]['audio'] signal_len = batch[key]['audio_lens'] speaker_id = batch[key]['speaker_id'] sv_logits, sv_emb, _, _, _ = self.forward(input_signal=signal, input_signal_length=signal_len) pred_speaker = torch.argmax(sv_logits, dim=1) sv_loss = self.sv_loss(logits=sv_logits, labels=speaker_id) loss += sv_loss if not self._cfg.combined_loss: optim_backbone.zero_grad() optim_downstream.zero_grad() self.manual_backward(sv_loss) optim_backbone.step() optim_downstream.step() correct = pred_speaker.eq(speaker_id.data.view_as(pred_speaker)).sum().item() acc = (correct / len(speaker_id)) * 100 self.log("t_sv_loss", sv_loss.item()) self.log("t_sv_accuracy", acc) elif key == "content": content_loss = 0 signal = batch[key]['audio'] signal_len = batch[key]['audio_lens'] target = batch[key]['text'] # (B, T) target_len = batch[key]['text_lens'] _, _, content_embedding, content_log_probs, encoded_len = self.forward( input_signal=signal, input_signal_length=signal_len ) ctc_loss = self.ctc_loss(content_log_probs, target, encoded_len, target_len) # check if ctc loss is nan if torch.isfinite(ctc_loss): self.log("t_ctc_loss", ctc_loss.item()) content_loss += ctc_loss else: logging.warning("ctc_loss is not finite") if self.pitch_augment: augmented_signal = batch[key]['audio_shifted'] if self.stop_gradient: with torch.no_grad(): _, _, content_embedding_aug, content_log_probs_aug, _ = self.forward( input_signal=augmented_signal, input_signal_length=signal_len ) else: _, _, content_embedding_aug, content_log_probs_aug, _ = self.forward( input_signal=augmented_signal, input_signal_length=signal_len ) if self.aug_loss_type == "mse": sim_loss = self.mse_loss(content_embedding, content_embedding_aug) elif self.aug_loss_type == "cosine": cosine_similarity = torch.nn.functional.cosine_similarity( content_embedding, content_embedding_aug, dim=-1 ).mean() sim_loss = 1.0 - cosine_similarity content_loss += self._cfg.augment_sim_alpha * sim_loss self.log("t_sim_loss", sim_loss.item()) if self.augment_ctc: ctc_loss_aug = self.ctc_loss(content_log_probs_aug, target, encoded_len, target_len) if torch.isfinite(ctc_loss_aug): content_loss += ctc_loss_aug self.log("t_ctc_loss_aug", ctc_loss_aug.item()) else: logging.warning("ctc_loss_aug is not finite. Add min duration to avoid getting here.") loss += content_loss if not self._cfg.combined_loss: optim_backbone.zero_grad() optim_downstream.zero_grad() self.manual_backward(content_loss) optim_backbone.step() optim_downstream.step() if isinstance(content_loss, torch.Tensor): self.log("t_content_loss", content_loss.item()) if self._cfg.combined_loss: optim_backbone.zero_grad() optim_downstream.zero_grad() self.manual_backward(loss) optim_backbone.step() optim_downstream.step() if schedulers is not None: sch1, sch2 = schedulers sch1.step() sch2.step() if self.trainer.global_step % 10 == 0: self.log("lr_backbone", optim_backbone.param_groups[0]['lr']) self.log("lr_downstream", optim_downstream.param_groups[0]['lr']) self.log("t_loss", loss) def validation_step(self, batch, batch_idx): loss_total = 0 for key in batch.keys(): if key == 'sv': signal = batch[key]['audio'] signal_len = batch[key]['audio_lens'] speaker_id = batch[key]['speaker_id'] sv_logits, sv_emb, _, _, _ = self.forward(input_signal=signal, input_signal_length=signal_len) pred_speaker = torch.argmax(sv_logits, dim=1) sv_loss = self.sv_loss(logits=sv_logits, labels=speaker_id) loss_total += sv_loss correct = pred_speaker.eq(speaker_id.data.view_as(pred_speaker)).sum().item() acc = (correct / len(speaker_id)) * 100 acc_val = torch.as_tensor(acc) if key == 'content': content_loss = 0 signal = batch[key]['audio'] signal_len = batch[key]['audio_lens'] target = batch[key]['text'] # (B, T) target_len = batch[key]['text_lens'] _, _, content_embedding, content_log_probs, encoded_len = self.forward( input_signal=signal, input_signal_length=signal_len ) ctc_loss = self.ctc_loss(content_log_probs, target, encoded_len, target_len) # check if ctc loss is nan if torch.isfinite(ctc_loss): content_loss += ctc_loss else: logging.warning("ctc_loss is not finite. Add min duration to avoid getting here.") if self.pitch_augment: augmented_signal = batch[key]['audio_shifted'] _, _, content_embedding_aug, content_log_probs_aug, _ = self.forward( input_signal=augmented_signal, input_signal_length=signal_len ) if self.aug_loss_type == "mse": sim_loss = self.mse_loss(content_embedding, content_embedding_aug) elif self.aug_loss_type == "cosine": cosine_similarity = torch.nn.functional.cosine_similarity( content_embedding, content_embedding_aug, dim=-1 ).mean() sim_loss = 1.0 - cosine_similarity content_loss += self._cfg.augment_sim_alpha * sim_loss loss_total += content_loss cers = [] for _idx in range(target.shape[0]): item_log_prob = content_log_probs[:, _idx, :][: encoded_len[_idx]].cpu() item_target = target[_idx][: target_len[_idx]].cpu() _, predicted_str = self.ctc_decoder(item_log_prob) tokenizer = self._text_tokenizer target_str = tokenizer.sep.join(tokenizer._id2token[t] for t in item_target.tolist()) ed = edit_distance(list(predicted_str), list(target_str))['total'] if max(len(predicted_str), len(target_str)) > 0: normalized_ed = (1.0 * ed) / max(len(predicted_str), len(target_str)) else: normalized_ed = 1.0 cers.append(normalized_ed) return { 'val_loss': loss_total.cpu(), 'sv_loss': sv_loss.cpu(), 'ctc_loss': ctc_loss.cpu(), 'content_loss': content_loss.cpu(), 'accuracy_sv': acc_val.cpu(), 'cer': torch.tensor(cers).mean().cpu(), } def on_validation_epoch_end(self, outputs): collect = lambda key: torch.stack([x[key] for x in outputs if torch.isfinite(x[key])]).mean() val_loss = collect("val_loss") val_sv_loss = collect("sv_loss") val_ctc_loss = collect("ctc_loss") val_content_loss = collect("content_loss") accuracy_sv = collect("accuracy_sv") cer = collect("cer") self.log("val_loss", val_loss) self.log("sv_loss", val_sv_loss) self.log("val_ctc_loss", val_ctc_loss) self.log("val_content_loss", val_content_loss) self.log("accuracy_sv", accuracy_sv) self.log("cer", cer)