# Copyright (c) 2022, NVIDIA CORPORATION. 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. from math import ceil from typing import Any, Dict, List, Optional, Union import torch import torch.nn as nn from lightning.pytorch import Trainer from omegaconf import DictConfig from nemo.collections.asr.data import audio_to_text_dataset, ssl_dataset from nemo.collections.asr.data.audio_to_text_dali import DALIOutputs from nemo.collections.asr.data.audio_to_text_lhotse import LhotseSpeechToTextBpeDataset from nemo.collections.asr.modules.ssl_modules.masking import ConvFeatureMaksingWrapper from nemo.collections.asr.parts.mixins import ASRModuleMixin from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config from nemo.collections.common.data.utils import move_data_to_device from nemo.collections.common.parts.preprocessing.parsers import make_parser from nemo.core.classes import ModelPT from nemo.core.classes.common import PretrainedModelInfo, typecheck from nemo.core.classes.mixins import AccessMixin, set_access_cfg from nemo.core.neural_types import ( AcousticEncodedRepresentation, AudioSignal, LabelsType, LengthsType, LogprobsType, NeuralType, SpectrogramType, ) from nemo.utils import logging __all__ = ['SpeechEncDecSelfSupervisedModel', 'EncDecMaskedTokenPredModel', 'EncDecDenoiseMaskedTokenPredModel'] class SpeechEncDecSelfSupervisedModel(ModelPT, ASRModuleMixin, AccessMixin): """Base class for encoder-decoder models used for self-supervised encoder pre-training""" @classmethod def list_available_models(cls) -> List[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 def __init__(self, cfg: DictConfig, trainer: Trainer = None): # Get global rank and total number of GPU workers for IterableDataset partitioning, if applicable # Global_rank and local_rank is set by LightningModule in Lightning 1.2.0 self.world_size = 1 if trainer is not None: self.world_size = trainer.world_size super().__init__(cfg=cfg, trainer=trainer) self.preprocessor = SpeechEncDecSelfSupervisedModel.from_config_dict(self._cfg.preprocessor) self.encoder = SpeechEncDecSelfSupervisedModel.from_config_dict(self._cfg.encoder) self.decoder_losses = None if "loss_list" in self._cfg: self.decoder_losses = {} self.loss_alphas = {} self.start_step = {} self.output_from_layer = {} self.transpose_encoded = {} self.targets_from_loss = {} self.decoder_losses_active = {} # need to be separate for moduledict for decoder_loss_name, decoder_loss_cfg in self._cfg.loss_list.items(): if not decoder_loss_cfg.get("is_active", True): # active by default continue new_decoder_loss = { 'decoder': SpeechEncDecSelfSupervisedModel.from_config_dict(decoder_loss_cfg.decoder), 'loss': SpeechEncDecSelfSupervisedModel.from_config_dict(decoder_loss_cfg.loss), } new_decoder_loss = nn.ModuleDict(new_decoder_loss) self.decoder_losses[decoder_loss_name] = new_decoder_loss self.loss_alphas[decoder_loss_name] = decoder_loss_cfg.get("loss_alpha", 1.0) self.output_from_layer[decoder_loss_name] = decoder_loss_cfg.get("output_from_layer", None) self.targets_from_loss[decoder_loss_name] = decoder_loss_cfg.get("targets_from_loss", None) self.start_step[decoder_loss_name] = decoder_loss_cfg.get("start_step", 0) self.transpose_encoded[decoder_loss_name] = decoder_loss_cfg.get("transpose_encoded", False) self.decoder_losses_active[decoder_loss_name] = True self.decoder_losses = nn.ModuleDict(self.decoder_losses) else: self.decoder_ssl = SpeechEncDecSelfSupervisedModel.from_config_dict(self._cfg.decoder) self.loss = SpeechEncDecSelfSupervisedModel.from_config_dict(self._cfg.loss) self.spec_augmentation = SpeechEncDecSelfSupervisedModel.from_config_dict(self._cfg.spec_augment) # dropout for features/spectrograms (applied before masking) self.dropout_features = ( torch.nn.Dropout(self._cfg.dropout_features) if "dropout_features" in self._cfg else None ) # dropout for targets (applied before quantization) self.dropout_features_q = ( torch.nn.Dropout(self._cfg.dropout_features_q) if "dropout_features_q" in self._cfg else None ) # Feature penalty for preprocessor encodings (for Wav2Vec training) if "feature_penalty" in self._cfg: self.feat_pen, self.pen_factor = 0.0, self._cfg.feature_penalty else: self.feat_pen, self.pen_factor = None, None if "access" in self._cfg: set_access_cfg(self._cfg.access, self.model_guid) self.apply_masking = True def _setup_dataloader_from_config(self, config: Optional[Dict]): if 'augmentor' in config: augmentor = process_augmentations(config['augmentor']) else: augmentor = None # Automatically inject args from model config to dataloader config audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='sample_rate') if config.get("use_lhotse"): return get_lhotse_dataloader_from_config( config, global_rank=self.global_rank, world_size=self.world_size, dataset=LhotseSpeechToTextBpeDataset( tokenizer=make_parser( labels=config.get('labels', None), name=config.get('parser', 'en'), unk_id=config.get('unk_index', -1), blank_id=config.get('blank_index', -1), do_normalize=config.get('normalize_transcripts', False), ), ), ) shuffle = config['shuffle'] device = 'gpu' if torch.cuda.is_available() else 'cpu' if config.get('use_dali', False): device_id = self.local_rank if device == 'gpu' else None dataset = audio_to_text_dataset.get_dali_char_dataset( config=config, shuffle=shuffle, device_id=device_id, global_rank=self.global_rank, world_size=self.world_size, preprocessor_cfg=self._cfg.preprocessor, ) return dataset # Instantiate tarred dataset loader or normal dataset loader if config.get('is_tarred', False): if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or ( 'manifest_filepath' in config and config['manifest_filepath'] is None ): logging.warning( "Could not load dataset as `manifest_filepath` was None or " f"`tarred_audio_filepaths` is None. Provided config : {config}" ) return None shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0 dataset = audio_to_text_dataset.get_tarred_dataset( config=config, shuffle_n=shuffle_n, global_rank=self.global_rank, world_size=self.world_size, augmentor=augmentor, ) shuffle = False else: if 'manifest_filepath' in config and config['manifest_filepath'] is None: logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}") return None dataset = audio_to_text_dataset.get_char_dataset(config=config, augmentor=augmentor) if hasattr(dataset, 'collate_fn'): collate_fn = dataset.collate_fn elif hasattr(dataset.datasets[0], 'collate_fn'): # support datasets that are lists of entries collate_fn = dataset.datasets[0].collate_fn else: # support datasets that are lists of lists collate_fn = dataset.datasets[0].datasets[0].collate_fn return torch.utils.data.DataLoader( dataset=dataset, batch_size=config['batch_size'], collate_fn=collate_fn, drop_last=config.get('drop_last', False), shuffle=shuffle, num_workers=config.get('num_workers', 0), pin_memory=config.get('pin_memory', False), ) def setup_training_data(self, train_data_config: Optional[Union[DictConfig, Dict]]): """ Sets up the training data loader via a Dict-like object. Args: train_data_config: A config that contains the information regarding construction of an ASR Training dataset. Supported Datasets: - :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset` - :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset` - :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset` - :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset` - :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset` """ if 'shuffle' not in train_data_config: train_data_config['shuffle'] = True # preserve config self._update_dataset_config(dataset_name='train', config=train_data_config) self._train_dl = self._setup_dataloader_from_config(config=train_data_config) # Need to set this because if using an IterableDataset, the length of the dataloader is the total number # of samples rather than the number of batches, and this messes up the tqdm progress bar. # So we set the number of steps manually (to the correct number) to fix this. if ( self._train_dl is not None and hasattr(self._train_dl, 'dataset') and isinstance(self._train_dl.dataset, torch.utils.data.IterableDataset) ): # We also need to check if limit_train_batches is already set. # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches, # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0). if self._trainer is not None and isinstance(self._trainer.limit_train_batches, float): self._trainer.limit_train_batches = int( self._trainer.limit_train_batches * ceil((len(self._train_dl.dataset) / self.world_size) / train_data_config['batch_size']) ) elif self._trainer is None: logging.warning( "Model Trainer was not set before constructing the dataset, incorrect number of " "training batches will be used. Please set the trainer and rebuild the dataset." ) def setup_validation_data(self, val_data_config: Optional[Union[DictConfig, Dict]]): """ Sets up the validation data loader via a Dict-like object. Args: val_data_config: A config that contains the information regarding construction of an ASR Training dataset. Supported Datasets: - :class:`~nemo.collections.asr.data.audio_to_text.AudioToCharDataset` - :class:`~nemo.collections.asr.data.audio_to_text.AudioToBPEDataset` - :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset` - :class:`~nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset` - :class:`~nemo.collections.asr.data.audio_to_text_dali.AudioToCharDALIDataset` """ if 'shuffle' not in val_data_config: val_data_config['shuffle'] = False # preserve config self._update_dataset_config(dataset_name='validation', config=val_data_config) self._validation_dl = self._setup_dataloader_from_config(config=val_data_config) # Need to set this because if using an IterableDataset, the length of the dataloader is the total number # of samples rather than the number of batches, and this messes up the tqdm progress bar. # So we set the number of steps manually (to the correct number) to fix this. if ( self._validation_dl is not None and hasattr(self._validation_dl, 'dataset') and isinstance(self._validation_dl.dataset, torch.utils.data.IterableDataset) ): # We also need to check if limit_train_batches is already set. # If it's an int, we assume that the user has set it to something sane, i.e. <= # training batches, # and don't change it. Otherwise, adjust batches accordingly if it's a float (including 1.0). if isinstance(self._trainer.limit_val_batches, float): self._trainer.limit_val_batches = int( self._trainer.limit_val_batches * ceil((len(self._validation_dl.dataset) / self.world_size) / val_data_config['batch_size']) ) @property def input_types(self) -> Optional[Dict[str, NeuralType]]: if hasattr(self.preprocessor, '_sample_rate'): input_signal_eltype = AudioSignal(freq=self.preprocessor._sample_rate) else: input_signal_eltype = AudioSignal() return { "input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True), "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True), "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "targets": NeuralType(('B', 'T'), LabelsType(), optional=True), "target_lengths": NeuralType(tuple('B'), LengthsType(), optional=True), } @property def output_types(self) -> Optional[Dict[str, NeuralType]]: return { "spectrograms": NeuralType(('B', 'D', 'T'), SpectrogramType()), "spec_masks": NeuralType(('B', 'D', 'T'), SpectrogramType()), "encoded": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()), "encoded_len": NeuralType(tuple('B'), LengthsType()), } @typecheck() def forward( self, input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None, ): """ Forward pass of the model. Args: input_signal: Tensor that represents a batch of raw audio signals, of shape [B, T]. T here represents timesteps, with 1 second of audio represented as `self.sample_rate` number of floating point values. input_signal_length: Vector of length B, that contains the individual lengths of the audio sequences. processed_signal: Tensor that represents a batch of processed audio signals, of shape (B, D, T) that has undergone processing via some DALI preprocessor. processed_signal_length: Vector of length B, that contains the individual lengths of the processed audio sequences. Returns: A tuple of 4 elements - 1) Processed spectrograms of shape [B, D, T]. 2) Masks applied to spectrograms of shape [B, D, T]. 3) The encoded features tensor of shape [B, D, T]. 2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B]. """ # Reset access registry if self.is_access_enabled(self.model_guid): self.reset_registry() # Check for special flag for validation step if hasattr(self, '_in_validation_step'): in_validation_step = self._in_validation_step else: in_validation_step = False # reset module registry from AccessMixin if ( (self.training or in_validation_step) and self.decoder_losses is not None and self.output_from_layer is not None and len(self.output_from_layer) > 0 ): layer_names = list(self.output_from_layer.values()) register_layer = any([name is not None for name in layer_names]) if register_layer: self.access_cfg['save_encoder_tensors'] = True self.set_access_enabled(access_enabled=True, guid=self.model_guid) has_input_signal = input_signal is not None and input_signal_length is not None has_processed_signal = processed_signal is not None and processed_signal_length is not None if (has_input_signal ^ has_processed_signal) == False: raise ValueError( f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive " " with ``processed_signal`` and ``processed_signal_len`` arguments." ) if not has_processed_signal: processed_signal, processed_signal_length = self.preprocessor( input_signal=input_signal, length=input_signal_length, ) if self.pen_factor: self.feat_pen = processed_signal.float().pow(2).mean() * self.pen_factor spectrograms = processed_signal.detach().clone() if self.dropout_features: processed_signal = self.dropout_features(processed_signal) if self.dropout_features_q: spectrograms = self.dropout_features_q(spectrograms) if self.apply_masking: processed_signal = self.spec_augmentation(input_spec=processed_signal, length=processed_signal_length) masked_spectrograms = processed_signal.detach() spec_masks = torch.logical_and(masked_spectrograms < 1e-5, masked_spectrograms > -1e-5).float() for idx, proc_len in enumerate(processed_signal_length): spec_masks[idx, :, proc_len:] = 0.0 encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length) return spectrograms, spec_masks, encoded, encoded_len def decoder_loss_step(self, spectrograms, spec_masks, encoded, encoded_len, targets=None, target_lengths=None): """ Forward pass through all decoders and calculate corresponding losses. Args: spectrograms: Processed spectrograms of shape [B, D, T]. spec_masks: Masks applied to spectrograms of shape [B, D, T]. encoded: The encoded features tensor of shape [B, D, T]. encoded_len: The lengths of the acoustic sequence after propagation through the encoder, of shape [B]. targets: Optional target labels of shape [B, T] target_lengths: Optional target label lengths of shape [B] Returns: A tuple of 2 elements - 1) Total sum of losses weighted by corresponding loss_alphas 2) Dictionary of unweighted losses """ loss_val_dict = {} if self.decoder_losses is None: if hasattr(self.decoder_ssl, "needs_labels") and self.decoder_ssl.needs_labels: outputs = self.decoder_ssl(encoder_output=encoded, targets=targets, target_lengths=target_lengths) else: outputs = self.decoder_ssl(encoder_output=encoded) if self.loss.needs_labels: loss_value = self.loss( spec_masks=spec_masks, decoder_outputs=outputs, targets=targets, decoder_lengths=encoded_len, target_lengths=target_lengths, ) else: loss_value = self.loss(spectrograms=spectrograms, spec_masks=spec_masks, decoder_outputs=outputs) else: loss_value = encoded.new_zeros(1) outputs = {} registry = self.get_module_registry(self.encoder) for dec_loss_name, dec_loss in self.decoder_losses.items(): # loop through decoders and corresponding losses if not self.decoder_losses_active[dec_loss_name]: continue if self.output_from_layer[dec_loss_name] is None: dec_input = encoded else: # extract output from specified layer using AccessMixin registry dec_input = registry[self.output_from_layer[dec_loss_name]]['encoder'][-1] if self.transpose_encoded[dec_loss_name]: dec_input = dec_input.transpose(-2, -1) if self.targets_from_loss[dec_loss_name] is not None: # extract targets from specified loss target_loss = self.targets_from_loss[dec_loss_name] targets = self.decoder_losses[target_loss]['loss'].target_ids target_lengths = self.decoder_losses[target_loss]['loss'].target_lengths if target_lengths is None: target_lengths = encoded_len if hasattr(dec_loss['decoder'], "needs_labels") and dec_loss['decoder'].needs_labels: # if we are using a decoder which needs labels, provide them outputs[dec_loss_name] = dec_loss['decoder']( encoder_output=dec_input, targets=targets, target_lengths=target_lengths ) else: outputs[dec_loss_name] = dec_loss['decoder'](encoder_output=dec_input) current_loss = dec_loss['loss'] if current_loss.needs_labels: # if we are using a loss which needs labels, provide them current_loss_value = current_loss( spec_masks=spec_masks, decoder_outputs=outputs[dec_loss_name], targets=targets, decoder_lengths=encoded_len, target_lengths=target_lengths, ) else: current_loss_value = current_loss( spectrograms=spectrograms, spec_masks=spec_masks, decoder_outputs=outputs[dec_loss_name], decoder_lengths=encoded_len, ) loss_value = loss_value + current_loss_value * self.loss_alphas[dec_loss_name] loss_val_dict[dec_loss_name] = current_loss_value return loss_value, loss_val_dict # PTL-specific methods def training_step(self, batch, batch_nb): signal, signal_len, targets, target_lengths = batch if isinstance(batch, DALIOutputs) and batch.has_processed_signal: spectrograms, spec_masks, encoded, encoded_len = self.forward( processed_signal=signal, processed_signal_length=signal_len, ) else: spectrograms, spec_masks, encoded, encoded_len = self.forward( input_signal=signal, input_signal_length=signal_len, ) if self.decoder_losses is not None: for dec_loss_name, dec_loss in self.decoder_losses.items(): self.decoder_losses_active[dec_loss_name] = self.trainer.global_step >= self.start_step[dec_loss_name] loss = dec_loss['loss'] if hasattr(loss, "set_num_updates"): loss.set_num_updates(self.trainer.global_step) else: if hasattr(self.loss, "set_num_updates"): self.loss.set_num_updates(self.trainer.global_step) loss_value, loss_val_dict = self.decoder_loss_step( spectrograms, spec_masks, encoded, encoded_len, targets, target_lengths ) tensorboard_logs = { 'learning_rate': self._optimizer.param_groups[0]['lr'], 'global_step': self.trainer.global_step, } for loss_name, loss_val in loss_val_dict.items(): tensorboard_logs['train_' + loss_name] = loss_val if self.feat_pen: loss_value += self.feat_pen # Reset access registry self.reset_registry() return {'loss': loss_value, 'log': tensorboard_logs} def validation_pass(self, batch, batch_idx, dataloader_idx=0): # Set flag to register tensors self._in_validation_step = True signal, signal_len, targets, target_lengths = batch if isinstance(batch, DALIOutputs) and batch.has_processed_signal: spectrograms, spec_masks, encoded, encoded_len = self.forward( processed_signal=signal, processed_signal_length=signal_len, ) else: spectrograms, spec_masks, encoded, encoded_len = self.forward( input_signal=signal, input_signal_length=signal_len, ) if self.decoder_losses is not None: for dec_loss_name, dec_loss in self.decoder_losses.items(): self.decoder_losses_active[dec_loss_name] = self.trainer.global_step >= self.start_step[dec_loss_name] loss_value, _ = self.decoder_loss_step(spectrograms, spec_masks, encoded, encoded_len, targets, target_lengths) if self.feat_pen: loss_value += self.feat_pen # reset access registry self.reset_registry() del self._in_validation_step metrics = {'val_loss': loss_value} return metrics def validation_step(self, batch, batch_idx, dataloader_idx=0): metrics = self.validation_pass(batch, batch_idx, dataloader_idx) if type(self.trainer.val_dataloaders) == list and len(self.trainer.val_dataloaders) > 1: self.validation_step_outputs[dataloader_idx].append(metrics) else: self.validation_step_outputs.append(metrics) return metrics def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0): val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() tensorboard_logs = {'val_loss': val_loss_mean} return {'val_loss': val_loss_mean, 'log': tensorboard_logs} class EncDecMaskedTokenPredModel(SpeechEncDecSelfSupervisedModel): """ Speech self-supervised model that performs masked token prediction on the encoder output. """ def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any: """ PTL hook: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#transfer-batch-to-device """ batch = move_data_to_device(batch, device) return batch @classmethod def list_available_models(cls) -> List[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 = [] return results def __init__(self, cfg: DictConfig, trainer: Trainer = None): super().__init__(cfg, trainer) del self.decoder_ssl # delete unused decoder from parent class if self.cfg.get("mask_position", "pre_conv") == "post_conv": # adjust config for post-convolution masking self.cfg.quantizer.feat_in = self.cfg.encoder.d_model self.cfg.masking.feat_in = self.cfg.encoder.d_model self.cfg.masking.block_size = self.cfg.masking.block_size // self.cfg.encoder.subsampling_factor self.cfg.loss.combine_time_steps = 1 self.quantizer = self.from_config_dict(self.cfg.quantizer) self.mask_processor = self.from_config_dict(self.cfg.masking) self.encoder = self.from_config_dict(self.cfg.encoder) self.decoder = self.from_config_dict(self.cfg.decoder) self.loss = self.from_config_dict(self.cfg.loss) self.pre_encoder = None if self.cfg.get("mask_position", "pre_conv") == "post_conv": # hacked to mask features after convolutional sub-sampling self.pre_encoder = ConvFeatureMaksingWrapper(self.encoder.pre_encode, self.mask_processor) self.encoder.pre_encode = self.pre_encoder @property def oomptimizer_schema(self) -> dict: """ Return a typing schema for optimal batch size calibration for various sequence lengths using OOMptimizer. """ return { "cls": tuple, "inputs": [ {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"}, {"type": NeuralType(("B",), LengthsType()), "seq_length": "input"}, ], } @property def input_types(self) -> Optional[Dict[str, NeuralType]]: if hasattr(self.preprocessor, '_sample_rate'): input_signal_eltype = AudioSignal(freq=self.preprocessor._sample_rate) else: input_signal_eltype = AudioSignal() return { "input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True), "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True), "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "apply_mask": NeuralType(optional=True), } @property def output_types(self) -> Optional[Dict[str, NeuralType]]: if self.cfg.num_books == 1 and self.cfg.squeeze_single: logprobs = NeuralType(('B', 'T', 'C'), LogprobsType()) tokens = NeuralType(('B', 'T'), LabelsType()) else: logprobs = NeuralType(('B', 'T', 'C', 'H'), LogprobsType()) tokens = NeuralType(('B', 'T', 'H'), LabelsType()) return { "logprobs": logprobs, "encoded_len": NeuralType(tuple('B'), LengthsType()), "masks": NeuralType(('B', 'D', 'T'), SpectrogramType()), "tokens": tokens, } @typecheck() def forward( self, input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None, apply_mask=False, ): has_input_signal = input_signal is not None and input_signal_length is not None has_processed_signal = processed_signal is not None and processed_signal_length is not None if (has_input_signal ^ has_processed_signal) == False: raise ValueError( f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive " " with ``processed_signal`` and ``processed_signal_len`` arguments." ) if not has_processed_signal: processed_signal, processed_signal_length = self.preprocessor( input_signal=input_signal, length=input_signal_length, ) if self.pre_encoder is not None: # mask after convolutional sub-sampling self.pre_encoder.set_masking_enabled(apply_mask=apply_mask) encoded, encoded_len = self.encoder(audio_signal=processed_signal, length=processed_signal_length) masks = self.pre_encoder.get_current_mask() feats = self.pre_encoder.get_current_feat() _, tokens = self.quantizer(input_signal=feats.transpose(1, 2)) else: _, tokens = self.quantizer(input_signal=processed_signal) if apply_mask: masked_signal, masks = self.mask_processor( input_feats=processed_signal, input_lengths=processed_signal_length ) else: masked_signal = processed_signal masks = torch.zeros_like(processed_signal) encoded, encoded_len = self.encoder(audio_signal=masked_signal, length=processed_signal_length) log_probs = self.decoder(encoder_output=encoded) return log_probs, encoded_len, masks, tokens def training_step(self, batch, batch_idx=0): input_signal, input_signal_length = batch[0], batch[1] if isinstance(batch, DALIOutputs) and batch.has_processed_signal: log_probs, encoded_len, masks, tokens = self.forward( processed_signal=input_signal, processed_signal_length=input_signal_length, apply_mask=True ) else: log_probs, encoded_len, masks, tokens = self.forward( input_signal=input_signal, input_signal_length=input_signal_length, apply_mask=True ) loss_value = self.loss(masks=masks, decoder_outputs=log_probs, targets=tokens, decoder_lengths=encoded_len) tensorboard_logs = { 'learning_rate': self._optimizer.param_groups[0]['lr'], 'global_step': self.trainer.global_step, 'train_loss': loss_value, } return {'loss': loss_value, 'log': tensorboard_logs} def inference_pass(self, batch, batch_idx=0, dataloader_idx=0, mode='val', apply_mask=False): input_signal, input_signal_length = batch[0], batch[1] if isinstance(batch, DALIOutputs) and batch.has_processed_signal: log_probs, encoded_len, masks, tokens = self.forward( processed_signal=input_signal, processed_signal_length=input_signal_length, apply_mask=apply_mask ) else: log_probs, encoded_len, masks, tokens = self.forward( input_signal=input_signal, input_signal_length=input_signal_length, apply_mask=apply_mask ) loss_value = self.loss(masks=masks, decoder_outputs=log_probs, targets=tokens, decoder_lengths=encoded_len) return {f'{mode}_loss': loss_value} def validation_step(self, batch, batch_idx=0, dataloader_idx=0): metrics = self.inference_pass(batch, batch_idx, dataloader_idx, apply_mask=True) if type(self.trainer.val_dataloaders) == list and len(self.trainer.val_dataloaders) > 1: self.validation_step_outputs[dataloader_idx].append(metrics) else: self.validation_step_outputs.append(metrics) return metrics def test_step(self, batch, batch_idx=0, dataloader_idx=0): metrics = self.inference_pass(batch, batch_idx, dataloader_idx, mode="test", apply_mask=True) if type(self.trainer.val_dataloaders) == list and len(self.trainer.val_dataloaders) > 1: self.validation_step_outputs[dataloader_idx].append(metrics) else: self.validation_step_outputs.append(metrics) return metrics def multi_validation_epoch_end(self, outputs: list, dataloader_idx: int = 0): loss_list = [] for i, x in enumerate(outputs): if not isinstance(x, dict): logging.warning(f'Batch {i} output in validation dataloader {dataloader_idx} is not a dictionary: {x}') if 'val_loss' in x: loss_list.append(x['val_loss']) else: logging.warning( f'Batch {i} output in validation dataloader {dataloader_idx} does not have key `val_loss`: {x}' ) if len(loss_list) == 0: logging.warning( f'Epoch {self.current_epoch} received no batches for validation dataloader {dataloader_idx}.' ) return {} val_loss_mean = torch.stack(loss_list).mean() tensorboard_logs = {'val_loss': val_loss_mean} return {'val_loss': val_loss_mean, 'log': tensorboard_logs} def multi_test_epoch_end(self, outputs, dataloader_idx: int = 0): test_loss_mean = torch.stack([x['test_loss'] for x in outputs]).mean() tensorboard_logs = {'test_loss': test_loss_mean} return {'test_loss': test_loss_mean, 'log': tensorboard_logs} class EncDecDenoiseMaskedTokenPredModel(EncDecMaskedTokenPredModel): """ Model class that performs denoising and masked token prediction for speech self-supervised learning. Please refer to the NEST paper for more details: https://arxiv.org/abs/2408.13106 """ @property def oomptimizer_schema(self) -> dict: """ Return a typing schema for optimal batch size calibration for various sequence lengths using OOMptimizer. """ return { "cls": ssl_dataset.AudioNoiseBatch, "inputs": [ {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input", "name": "audio"}, {"type": NeuralType(("B",), LengthsType()), "seq_length": "input", "name": "audio_len"}, {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input", "name": "noise"}, {"type": NeuralType(("B",), LengthsType()), "seq_length": "input", "name": "noise_len"}, {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input", "name": "noisy_audio"}, {"type": NeuralType(("B",), LengthsType()), "seq_length": "input", "name": "noisy_audio_len"}, ], } def __init__(self, cfg: DictConfig, trainer: Trainer = None): super().__init__(cfg, trainer) def _setup_dataloader_from_config(self, config: Optional[Dict]): audio_to_text_dataset.inject_dataloader_value_from_model_config(self.cfg, config, key='sample_rate') if config.get("use_lhotse"): return get_lhotse_dataloader_from_config( config, global_rank=self.global_rank, world_size=self.world_size, dataset=ssl_dataset.LhotseAudioNoiseDataset( noise_manifest=config.get('noise_manifest', None), batch_augmentor_cfg=config.get('batch_augmentor', None), ), ) dataset = ssl_dataset.get_audio_noise_dataset_from_config( config, global_rank=self.global_rank, world_size=self.world_size, ) shuffle = config['shuffle'] if isinstance(dataset, torch.utils.data.IterableDataset): shuffle = False if hasattr(dataset, 'collate_fn'): collate_fn = dataset.collate_fn elif hasattr(dataset.datasets[0], 'collate_fn'): # support datasets that are lists of entries collate_fn = dataset.datasets[0].collate_fn else: # support datasets that are lists of lists collate_fn = dataset.datasets[0].datasets[0].collate_fn return torch.utils.data.DataLoader( dataset=dataset, batch_size=config['batch_size'], collate_fn=collate_fn, drop_last=config.get('drop_last', False), shuffle=shuffle, num_workers=config.get('num_workers', 0), pin_memory=config.get('pin_memory', False), ) @property def input_types(self) -> Optional[Dict[str, NeuralType]]: if hasattr(self.preprocessor, '_sample_rate'): input_signal_eltype = AudioSignal(freq=self.preprocessor._sample_rate) else: input_signal_eltype = AudioSignal() return { "input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True), "input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "processed_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True), "processed_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "noise_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True), "noise_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "processed_noise_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True), "processed_noise_length": NeuralType(tuple('B'), LengthsType(), optional=True), "noisy_input_signal": NeuralType(('B', 'T'), input_signal_eltype, optional=True), "noisy_input_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "processed_noisy_signal": NeuralType(('B', 'D', 'T'), SpectrogramType(), optional=True), "processed_noisy_signal_length": NeuralType(tuple('B'), LengthsType(), optional=True), "apply_mask": NeuralType(optional=True), } @property def output_types(self) -> Optional[Dict[str, NeuralType]]: if self.cfg.num_books == 1 and self.cfg.squeeze_single: logprobs = NeuralType(('B', 'T', 'C'), LogprobsType()) tokens = NeuralType(('B', 'T'), LabelsType()) else: logprobs = NeuralType(('B', 'T', 'C', 'H'), LogprobsType()) tokens = NeuralType(('B', 'T', 'H'), LabelsType()) return { "logprobs": logprobs, "encoded_len": NeuralType(tuple('B'), LengthsType()), "masks": NeuralType(('B', 'D', 'T'), SpectrogramType()), "tokens": tokens, } @typecheck() def forward( self, input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None, noise_signal=None, # noqa noise_signal_length=None, # noqa processed_noise_signal=None, # noqa processed_noise_signal_length=None, # noqa noisy_input_signal=None, noisy_input_signal_length=None, processed_noisy_input_signal=None, processed_noisy_input_signal_length=None, apply_mask=False, ): has_input_signal = input_signal is not None and input_signal_length is not None has_processed_signal = processed_signal is not None and processed_signal_length is not None if (has_input_signal ^ has_processed_signal) == False: raise ValueError( f"{self} Arguments ``input_signal`` and ``input_signal_length`` are mutually exclusive " " with ``processed_signal`` and ``processed_signal_len`` arguments." ) if not has_processed_signal: processed_signal, processed_signal_length = self.preprocessor( input_signal=input_signal, length=input_signal_length, ) ### Following code snipet is not used but kept for future reference # # has_noise_signal = noise_signal is not None and noise_signal_length is not None # has_processed_noise_signal = processed_noise_signal is not None and processed_noise_signal_length is not None # if (has_noise_signal ^ has_processed_noise_signal) == False: # raise ValueError( # f"{self} Arguments ``noise_signal`` and ``noise_signal_length`` are mutually exclusive " # " with ``processed_noise_signal`` and ``processed_noise_signal_len`` arguments." # ) # if not has_processed_noise_signal: # processed_noise_signal, processed_noise_signal_length = self.preprocessor( # input_signal=noise_signal, # length=noise_signal_length, # ) has_noisy_input_signal = noisy_input_signal is not None and noisy_input_signal_length is not None has_processed_noisy_input_signal = ( processed_noisy_input_signal is not None and processed_noisy_input_signal_length is not None ) if (has_noisy_input_signal ^ has_processed_noisy_input_signal) == False: raise ValueError( f"{self} Arguments ``noisy_input_signal`` and ``noisy_input_signal_length`` are mutually exclusive " " with ``processed_noisy_input_signal`` and ``processed_noisy_input_signal_len`` arguments." ) if not has_processed_noisy_input_signal: processed_noisy_input_signal, processed_noisy_input_signal_length = self.preprocessor( input_signal=noisy_input_signal, length=noisy_input_signal_length, ) if self.pre_encoder is not None: # mask after convolutional sub-sampling feats, _ = self.pre_encoder.pre_encode(x=processed_signal, lengths=processed_signal_length) _, tokens = self.quantizer(input_signal=feats.transpose(1, 2)) self.pre_encoder.set_masking_enabled(apply_mask=apply_mask) encoded, encoded_len = self.encoder( audio_signal=processed_noisy_input_signal, length=processed_noisy_input_signal_length ) masks = self.pre_encoder.get_current_mask() else: _, tokens = self.quantizer(input_signal=processed_signal) if apply_mask: masked_signal, masks = self.mask_processor( input_feats=processed_noisy_input_signal, input_lengths=processed_noisy_input_signal_length ) else: masked_signal = processed_noisy_input_signal masks = torch.zeros_like(processed_noisy_input_signal) encoded, encoded_len = self.encoder(audio_signal=masked_signal, length=processed_noisy_input_signal_length) log_probs = self.decoder(encoder_output=encoded) return log_probs, encoded_len, masks, tokens def training_step(self, batch: ssl_dataset.AudioNoiseBatch, batch_idx: int): log_probs, encoded_len, masks, tokens = self.forward( input_signal=batch.audio, input_signal_length=batch.audio_len, noise_signal=batch.noise, noise_signal_length=batch.noise_len, noisy_input_signal=batch.noisy_audio, noisy_input_signal_length=batch.noisy_audio_len, apply_mask=True, ) loss_value = self.loss(masks=masks, decoder_outputs=log_probs, targets=tokens, decoder_lengths=encoded_len) tensorboard_logs = { 'learning_rate': self._optimizer.param_groups[0]['lr'], 'global_step': self.trainer.global_step, 'train_loss': loss_value, } return {'loss': loss_value, 'log': tensorboard_logs} def inference_pass( self, batch: ssl_dataset.AudioNoiseBatch, batch_idx: int, dataloader_idx: int = 0, mode: str = 'val', apply_mask: bool = True, ): log_probs, encoded_len, masks, tokens = self.forward( input_signal=batch.audio, input_signal_length=batch.audio_len, noise_signal=batch.noise, noise_signal_length=batch.noise_len, noisy_input_signal=batch.noisy_audio, noisy_input_signal_length=batch.noisy_audio_len, apply_mask=apply_mask, ) loss_value = self.loss(masks=masks, decoder_outputs=log_probs, targets=tokens, decoder_lengths=encoded_len) return {f'{mode}_loss': loss_value}