nvidia-nemo--speech/nemo/collections/tts/models/ssl_tts.py

# 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)