jone/Music_Source_Separation

1import math2from typing import Callable3 4import torch5import torch.nn as nn6from torchlibrosa.stft import STFT7 8from bytesep.models.pytorch_modules import Base9 10 11def l1(output: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor:12    r"""L1 loss.13 14    Args:15        output: torch.Tensor16        target: torch.Tensor17 18    Returns:19        loss: torch.float20    """21    return torch.mean(torch.abs(output - target))22 23 24def l1_wav(output: torch.Tensor, target: torch.Tensor, **kwargs) -> torch.Tensor:25    r"""L1 loss in the time-domain.26 27    Args:28        output: torch.Tensor29        target: torch.Tensor30 31    Returns:32        loss: torch.float33    """34    return l1(output, target)35 36 37class L1_Wav_L1_Sp(nn.Module, Base):38    def __init__(self):39        r"""L1 loss in the time-domain and L1 loss on the spectrogram."""40        super(L1_Wav_L1_Sp, self).__init__()41 42        self.window_size = 204843        hop_size = 44144        center = True45        pad_mode = "reflect"46        window = "hann"47 48        self.stft = STFT(49            n_fft=self.window_size,50            hop_length=hop_size,51            win_length=self.window_size,52            window=window,53            center=center,54            pad_mode=pad_mode,55            freeze_parameters=True,56        )57 58    def __call__(59        self, output: torch.Tensor, target: torch.Tensor, **kwargs60    ) -> torch.Tensor:61        r"""L1 loss in the time-domain and on the spectrogram.62 63        Args:64            output: torch.Tensor65            target: torch.Tensor66 67        Returns:68            loss: torch.float69        """70 71        # L1 loss in the time-domain.72        wav_loss = l1_wav(output, target)73 74        # L1 loss on the spectrogram.75        sp_loss = l1(76            self.wav_to_spectrogram(output, eps=1e-8),77            self.wav_to_spectrogram(target, eps=1e-8),78        )79 80        # sp_loss /= math.sqrt(self.window_size)81        # sp_loss *= 1.82 83        # Total loss.84        return wav_loss + sp_loss85 86        return sp_loss87 88 89def get_loss_function(loss_type: str) -> Callable:90    r"""Get loss function.91 92    Args:93        loss_type: str94 95    Returns:96        loss function: Callable97    """98 99    if loss_type == "l1_wav":100        return l1_wav101 102    elif loss_type == "l1_wav_l1_sp":103        return L1_Wav_L1_Sp()104 105    else:106        raise NotImplementedError107