def _setup_metrics(self): self._mse = MeanSquaredError() self._mae = MeanAbsoluteError()
class EncDecRegressionModel(_EncDecBaseModel): """Encoder decoder class for speech regression models. Model class creates training, validation methods for setting up data performing model forward pass. """ @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. """ result = [] return result def __init__(self, cfg: DictConfig, trainer: Trainer = None): if not cfg.get('is_regression_task', False): raise ValueError( f"EndDecRegressionModel requires the flag is_regression_task to be set as true" ) super().__init__(cfg=cfg, trainer=trainer) def _setup_preprocessor(self): return EncDecRegressionModel.from_config_dict(self._cfg.preprocessor) def _setup_encoder(self): return EncDecRegressionModel.from_config_dict(self._cfg.encoder) def _setup_decoder(self): return EncDecRegressionModel.from_config_dict(self._cfg.decoder) def _setup_loss(self): return MSELoss() def _setup_metrics(self): self._mse = MeanSquaredError() self._mae = MeanAbsoluteError() @property def output_types(self) -> Optional[Dict[str, NeuralType]]: return {"preds": NeuralType(tuple('B'), RegressionValuesType())} @typecheck() def forward(self, input_signal, input_signal_length): logits = super().forward(input_signal=input_signal, input_signal_length=input_signal_length) return logits.view(-1) # PTL-specific methods def training_step(self, batch, batch_idx): self.training_step_end() audio_signal, audio_signal_len, targets, targets_len = batch logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len) loss = self.loss(preds=logits, labels=targets) train_mse = self._mse(preds=logits, target=targets) train_mae = self._mae(preds=logits, target=targets) tensorboard_logs = { 'train_loss': loss, 'train_mse': train_mse, 'train_mae': train_mae, 'learning_rate': self._optimizer.param_groups[0]['lr'], } self.log_dict(tensorboard_logs) return {'loss': loss} def validation_step(self, batch, batch_idx, dataloader_idx: int = 0): audio_signal, audio_signal_len, targets, targets_len = batch logits = self.forward(input_signal=audio_signal, input_signal_length=audio_signal_len) loss_value = self.loss(preds=logits, labels=targets) val_mse = self._mse(preds=logits, target=targets) val_mae = self._mae(preds=logits, target=targets) return {'val_loss': loss_value, 'val_mse': val_mse, 'val_mae': val_mae} def test_step(self, batch, batch_idx, dataloader_idx: int = 0): logs = self.validation_step(batch, batch_idx, dataloader_idx) return { 'test_loss': logs['val_loss'], 'test_mse': logs['test_mse'], 'test_mae': logs['val_mae'] } def multi_validation_epoch_end(self, outputs, dataloader_idx: int = 0): val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean() val_mse = self._mse.compute() self._mse.reset() val_mae = self._mae.compute() self._mae.reset() tensorboard_logs = { 'val_loss': val_loss_mean, 'val_mse': val_mse, 'val_mae': val_mae } return { 'val_loss': val_loss_mean, 'val_mse': val_mse, 'val_mae': val_mae, '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() test_mse = self._mse.compute() self._mse.reset() test_mae = self._mae.compute() self._mae.reset() tensorboard_logs = { 'test_loss': test_loss_mean, 'test_mse': test_mse, 'test_mae': test_mae } return { 'test_loss': test_loss_mean, 'test_mse': test_mse, 'test_mae': test_mae, 'log': tensorboard_logs } @torch.no_grad() def transcribe(self, paths2audio_files: List[str], batch_size: int = 4) -> List[float]: """ Generate class labels for provided audio files. Use this method for debugging and prototyping. Args: paths2audio_files: (a list) of paths to audio files. \ Recommended length per file is approximately 1 second. batch_size: (int) batch size to use during inference. \ Bigger will result in better throughput performance but would use more memory. Returns: A list of predictions in the same order as paths2audio_files """ predictions = super().transcribe(paths2audio_files, batch_size, logprobs=True) return [float(pred) for pred in predictions] def _update_decoder_config(self, labels, cfg): OmegaConf.set_struct(cfg, False) if 'params' in cfg: cfg.params.num_classes = 1 else: cfg.num_classes = 1 OmegaConf.set_struct(cfg, True)