model/utils/abs_models.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Contains all abstract classes for speech or reverb models.

See `framework_details.md`.

"""

import os
import warnings
import torch
import numpy as np
from collections import OrderedDict
from lightning.pytorch import LightningModule
from torch import nn
from abc import ABC, abstractmethod
from model.utils.default_stft_istft import default_istft_module, default_stft_module
from model.utils.metrics import IgnoreLastSamplesMetricWrapper, DifferenceMetric
from model.utils.tensor_ops import crop_or_zero_pad_to_target_len
from model.utils.run_management import get_best_checkpoint
from enum import Enum
from lightning.pytorch.trainer.states import RunningStage


class SignalDomain(Enum):
    """Enum for signal domains used as input of joint loss."""

    TIME = 0
    STFT = 1
    # CRM=3


class ModelType(Enum):
    """Enum of the main components of the model.joint_model.JointModel."""

    DRY_SPEECH = "dry_speech_model"
    REVERB = "reverb_model"
    REREVERBERATION_LOSS = "rereverberation_loss"  # convolutive model and reverberation matching loss


class FirstLevelModule(LightningModule):
    """
    Handle metrics, logging, and STFT and ISTFT computation.

    See `framework_details.md`.
    """

    output_domain = SignalDomain.TIME  # default

    def __init__(self, metrics: list[nn.Module] = []):
        super().__init__()
        self.metrics = nn.ModuleList(metrics)
        self.stft_module = default_stft_module
        self.istft_module = default_istft_module
        self.fs = 16000
        self.metrics_all_batches = dict()

    def num_batches_for_current_trainer_stage(self, dataloader_idx=0):
        if self.trainer.state.stage == RunningStage.TRAINING:
            return self.trainer.num_training_batches
        if self.trainer.state.stage == RunningStage.VALIDATING:
            return self.trainer.num_val_batches[dataloader_idx]
        if self.trainer.state.stage == RunningStage.TESTING:
            return self.trainer.num_test_batches[dataloader_idx]
        if self.trainer.state.stage == RunningStage.SANITY_CHECKING:
            return self.trainer.num_sanity_val_batches[dataloader_idx]
        if self.trainer.state.stage == RunningStage.PREDICTING:
            return self.trainer.num_predict_batches[dataloader_idx]

    def log_for_future_aggregate(self, name, value, batch_idx, *args, dataloader_idx=0, **kwargs):
        if name not in self.metrics_all_batches or batch_idx == 0:
            size = (self.num_batches_for_current_trainer_stage(dataloader_idx=dataloader_idx),)
            self.metrics_all_batches[name] = torch.full(size=size, fill_value=torch.nan, device=self.device)
        self.metrics_all_batches[name][batch_idx] = value.detach()
        self.log(name, value, *args, **kwargs)

    def log_audios(self, waveform_batches_dict, scale=True):
        tensorboard = self.logger.experiment
        for k, v in waveform_batches_dict.items():
            wav = v[(0,) * (v.ndim - 1)]
            if wav.numel() > 0:
                if scale:
                    wav = wav / abs(wav).max()
                tensorboard.add_audio(k, wav, global_step=self.global_step, sample_rate=self.fs)
        tensorboard.close()

    def log_metrics_and_audios(self, pred_time_domain, target, batch_idx, dataloader_idx=0, input=None):
        for metric in self.metrics:
            metric_name = (
                "val_"
                + self.MODEL_TYPE.value
                + "_"
                + (
                    str(metric)
                    if isinstance(metric, (IgnoreLastSamplesMetricWrapper, DifferenceMetric))
                    else type(metric).__name__
                )
            )
            try:
                self.log_for_future_aggregate(
                    metric_name, metric(pred_time_domain, target), batch_idx=batch_idx, on_epoch=True
                )
                if self.MODEL_TYPE == ModelType.DRY_SPEECH and input is not None:
                    # Metric between input (wet speech) and target (dry speech)
                    self.log_for_future_aggregate(
                        metric_name + "_input", metric(input, target), batch_idx=batch_idx, on_epoch=True
                    )
            except:
                pass
        # log audios
        if batch_idx == 0:
            self.log_audios(
                {
                    self.MODEL_TYPE.value + "_target": target if self.MODEL_TYPE == ModelType.DRY_SPEECH else target[0],
                    self.MODEL_TYPE.value + "_prediction": pred_time_domain,
                }
            )
        if batch_idx == self.num_batches_for_current_trainer_stage(dataloader_idx=dataloader_idx) - 1:
            self.aggregate_metrics()

    def aggregate_metrics(self):
        if self.logger.log_dir is None:
            logdir = "lightning_logs/debug"
            warnings.warn(f"No logger.log_dir found, setting logdir to {logdir}")
        else:
            logdir = self.logger.log_dir
        metrics_folder = os.path.join(logdir, f"latest_results")
        os.makedirs(metrics_folder, exist_ok=True)
        for k, metrics_tensor in self.metrics_all_batches.items():
            self.log(k + "_median", torch.median(metrics_tensor))
            self.log(k + "_std", torch.std(metrics_tensor))
            np.save(
                os.path.join(metrics_folder, "".join(ki for ki in k if (ki.isalnum() or ki == "_")) + ".npy"),
                metrics_tensor.detach().cpu().numpy(),
            )
        self.metrics_all_batches.clear()

    def log_loss(self, loss, batch_idx=None):
        if loss is not None:
            self.log_for_future_aggregate(
                self._current_fx_name + "_" + self.MODEL_TYPE.value + "_loss", loss, batch_idx=batch_idx, on_epoch=True
            )

    def on_validation_epoch_end(self):
        self.aggregate_metrics()

    def on_test_epoch_end(self):
        self.aggregate_metrics()

    def on_train_epoch_end(self):
        self.aggregate_metrics()


class AbsSpeechOrReverbModel(ABC, FirstLevelModule):
    """
    Abstract parent class for SpeechModel or ReverbModel to handle training_step.

    See `framework_details.md`.
    """

    # Default stft and istft
    def __init__(self, metrics: list[nn.Module] = [], crop_input_to_target: bool = False):
        super().__init__(metrics=metrics)
        self.crop_input_to_target = crop_input_to_target

    @abstractmethod
    def internal_loss(self, pred, target):
        raise NotImplementedError()

    @abstractmethod
    def forward(self, input):
        """
        Forward pass

        Takes wet signal as input and returns the prediction, and all temporary tensors which are used to compute the internal loss.

        See `framework_details.md`.

        Parameters
        ----------
        input : torch.Tensor
            time-domain wet signal.

        Returns
        -------
        pred :
            prediction and all temporary tensors used to compute the internal loss

        """
        pred = NotImplemented
        return pred

    def get_time(self, pred, **kwargs):
        """See `framework_details.md`."""
        return pred

    def get_stft(self, pred, **kwargs):
        """See `framework_details.md`."""
        return self.stft_module(pred)

    def _get_pred_in_domain(self, pred, domain, **kwargs):
        if domain == SignalDomain.STFT:
            return self.get_stft(pred, **kwargs)
        if domain == SignalDomain.TIME:
            return self.get_time(pred, **kwargs)

    def _get_pred_good_domain(self, pred, **kwargs):
        return self._get_pred_in_domain(pred, self.output_domain, **kwargs)

    def _get_pred_all_domains(self, pred, domains, **kwargs):
        return {domain: self._get_pred_in_domain(pred, domain, **kwargs) for domain in domains}

    def training_step(self, batch, batch_idx):
        """
        Perform training step.

        See `framework_details.md`.

        Parameters
        ----------
        batch : tuple
            - If speech model: (y, s)
            - If reverb model: (y, (h, rir properties which are returned by the dataset).
        batch_idx : int
            batch_idx.

        Returns
        -------
        dict
            - "pred" :
                - If speech model, estimated dry speech (in apropriate domain for joint_loss_module).
                - If reverb model, estimated rir (in apropriate domain for joint_loss_module).
            - "loss" : Internal loss
        """
        input, target = batch
        if self.crop_input_to_target:
            input = input[..., : target.size(-1)]
        pred = self.forward(input)
        loss = self.internal_loss(pred, target)
        self.log_loss(loss, batch_idx=batch_idx)
        pred_good_domain = self._get_pred_good_domain(
            pred, length=(target.size(-1) if self.MODEL_TYPE == ModelType.DRY_SPEECH else None)
        )
        return {"pred": pred_good_domain, "loss": loss}

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        """
        Perform validation step.

        See `framework_details.md`.

        Parameters
        ----------
        batch : tuple
            - If speech model: (y, s)
            - If reverb model: (y, (h, rir properties which are returned by the dataset).
        batch_idx : int
            batch_idx.

        Returns
        -------
        dict
            - "pred" :
                - If speech model, estimated dry speech (in apropriate domain for joint_loss_module).
                - If reverb model, estimated rir (in apropriate domain for joint_loss_module).
            - "loss" : Internal loss
        """
        input, target = batch
        if self.crop_input_to_target:
            input = input[..., : target.size(-1)]
        pred = self.forward(input)
        loss = self.internal_loss(pred, target)
        self.log_loss(loss, batch_idx=batch_idx)

        # Metrics are supposed to all be time domain
        pred_time_domain = self.get_time(
            pred, length=(target.size(-1) if self.MODEL_TYPE == ModelType.DRY_SPEECH else None)
        )
        if self.output_domain != SignalDomain.TIME:
            pred_good_domain = self._get_pred_good_domain(pred)
        else:
            pred_good_domain = pred_time_domain

        self.log_metrics_and_audios(
            pred_time_domain=pred_time_domain,
            target=target,
            input=input,
            batch_idx=batch_idx,
            dataloader_idx=dataloader_idx,
        )
        return {"pred": pred_good_domain, "loss": loss}

    def load_state_dict_from_joint_model(self, ckpt_path: str):
        """
        Loads the state-dict of a speech or a reverb model from the checkpoint path of a JointModel.

        Parameters
        ----------
        ckpt_path : str
            - If ckeckpoint path: loads from this specific checkpoint.
            - If directory: Seeks the best checkpoint  in the directory according to a logged monitor.
            The monitor mode depends on ModelType. For speech, monitor mode is max, for reverb, monitor mode is min.
            See `this doc <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#lightning.pytorch.callbacks.ModelCheckpoint>`_

        Returns
        -------
        None.

        """
        if os.path.isdir(ckpt_path):
            print("dir is given as ckpt_path, loading best checkpoint")
            if self.MODEL_TYPE == ModelType.DRY_SPEECH:
                monitor_mode = max
            else:
                monitor_mode = min
            ckpt_path = get_best_checkpoint(ckpt_path, monitor_mode=monitor_mode)
            print(f"Best_checkpoint found {ckpt_path}")
        # map_location="cpu": ZeroGPU has no GPU outside @spaces.GPU calls
        state_dict = torch.load(ckpt_path, weights_only=False, map_location="cpu")["state_dict"]
        new_dict = OrderedDict()
        for k, v in state_dict.items():
            if self.MODEL_TYPE == ModelType.DRY_SPEECH and k.startswith("speech_model."):
                new_dict[k[13:]] = v
            # Ignore other modules (reverb model and joint_loss_module)
            if self.MODEL_TYPE == ModelType.REVERB and k.startswith("reverb_model."):
                new_dict[k[13:]] = v
            # else don't append the keys to the state dict
        self.load_state_dict(new_dict, strict=False)


class AbsSpeechModel(AbsSpeechOrReverbModel):
    """
    Abstract parent class for Speech Model.

    See `framework_details.md`.
    """

    MODEL_TYPE = ModelType.DRY_SPEECH


class AbsReverbModel(AbsSpeechOrReverbModel):
    """
    Abstract parent class for Reverb Model.

    See `framework_details.md`.
    """

    MODEL_TYPE = ModelType.REVERB


class OracleModel(ABC, FirstLevelModule):
    @abstractmethod
    def training_step(self, batch, batch_idx): ...

    def validation_step(self, batch, batch_idx):
        return self.training_step(batch, batch_idx)

    def internal_loss(self, pred, target):
        return None

    def forward(self, input):
        raise RuntimeError("Oracle model is not supposed to be used at inference")


class OracleSpeechModel(OracleModel, AbsSpeechModel):
    """
    Oracle speech model.

    See `framework_details.md`.
    """

    def training_step(self, batch, batch_idx):
        input, target = batch
        pred = target
        pred_good_domain = self._get_pred_good_domain(
            pred, length=(target.size(-1) if self.MODEL_TYPE == ModelType.DRY_SPEECH else None)
        )
        return {"pred": pred_good_domain, "loss": None}


class OracleReverbModel(OracleModel, AbsReverbModel):
    """
    Oracle reverb model.

    See `framework_details.md`.
    """

    def __init__(self, metrics: list[nn.Module] = [], target_len: int = 32000, crop_input_to_target: bool = False):
        super().__init__(metrics=metrics)
        self.crop_input_to_target = crop_input_to_target
        self.target_len = target_len

    def training_step(self, batch, batch_idx):
        _, (target, _) = batch
        target_cropped = crop_or_zero_pad_to_target_len(target, self.target_len)
        return {"pred": target_cropped, "loss": None}


class OracleParametersReverbModel(OracleReverbModel):
    """
    Abstract parent for reverb model with oracle parameters.

    See `framework_details.md`.
    """

    def training_step(self, batch, batch_idx):
        y, (h, rir_properties) = batch
        h_hat = self.convert_rir(h, rir_properties)
        return {"pred": h_hat, "loss": None}


class AbsJointLossModule(ABC, FirstLevelModule):
    """
    Abstract parent for joint loss.

    See `framework_details.md`.
    """

    def __init__(self, metrics: list[nn.Module] = []):
        super().__init__(metrics=metrics)

    @property
    @abstractmethod
    def SPEECH_INPUT_DOMAIN():
        return SignalDomain.TIME

    @property
    @abstractmethod
    def REVERB_INPUT_DOMAIN():
        return SignalDomain.TIME

    @abstractmethod
    def forward(self, s_hat, h_hat, s, h, y):
        loss = NotImplemented
        return loss

    def training_step(self, s_hat, h_hat, s, h, y, batch_idx):
        loss = self.forward(s_hat, h_hat, s, h, y)
        self.log_loss(loss, batch_idx=batch_idx)
        if True or not loss.isfinite():  # Disabled check in order to get better perf (no CPU-GPU sync)
            return loss
        return None

    def validation_step(self, s_hat, h_hat, s, h, y, batch_idx):
        raise NotImplementedError()