| import torch |
| from torch import nn |
|
|
| from .. import AudioSignal |
|
|
|
|
| class L1Loss(nn.L1Loss): |
| """L1 Loss between AudioSignals. Defaults |
| to comparing ``audio_data``, but any |
| attribute of an AudioSignal can be used. |
| |
| Parameters |
| ---------- |
| attribute : str, optional |
| Attribute of signal to compare, defaults to ``audio_data``. |
| weight : float, optional |
| Weight of this loss, defaults to 1.0. |
| """ |
|
|
| def __init__(self, attribute: str = "audio_data", weight: float = 1.0, **kwargs): |
| self.attribute = attribute |
| self.weight = weight |
| super().__init__(**kwargs) |
|
|
| def forward(self, x: AudioSignal, y: AudioSignal): |
| """ |
| Parameters |
| ---------- |
| x : AudioSignal |
| Estimate AudioSignal |
| y : AudioSignal |
| Reference AudioSignal |
| |
| Returns |
| ------- |
| torch.Tensor |
| L1 loss between AudioSignal attributes. |
| """ |
| if isinstance(x, AudioSignal): |
| x = getattr(x, self.attribute) |
| y = getattr(y, self.attribute) |
| return super().forward(x, y) |
|
|
|
|
| class SISDRLoss(nn.Module): |
| """ |
| Computes the Scale-Invariant Source-to-Distortion Ratio between a batch |
| of estimated and reference audio signals or aligned features. |
| |
| Parameters |
| ---------- |
| scaling : int, optional |
| Whether to use scale-invariant (True) or |
| signal-to-noise ratio (False), by default True |
| reduction : str, optional |
| How to reduce across the batch (either 'mean', |
| 'sum', or none).], by default ' mean' |
| zero_mean : int, optional |
| Zero mean the references and estimates before |
| computing the loss, by default True |
| clip_min : int, optional |
| The minimum possible loss value. Helps network |
| to not focus on making already good examples better, by default None |
| weight : float, optional |
| Weight of this loss, defaults to 1.0. |
| """ |
|
|
| def __init__( |
| self, |
| scaling: int = True, |
| reduction: str = "mean", |
| zero_mean: int = True, |
| clip_min: int = None, |
| weight: float = 1.0, |
| ): |
| self.scaling = scaling |
| self.reduction = reduction |
| self.zero_mean = zero_mean |
| self.clip_min = clip_min |
| self.weight = weight |
| super().__init__() |
|
|
| def forward(self, x: AudioSignal, y: AudioSignal): |
| eps = 1e-8 |
| |
| if isinstance(x, AudioSignal): |
| references = x.audio_data |
| estimates = y.audio_data |
| else: |
| references = x |
| estimates = y |
|
|
| nb = references.shape[0] |
| references = references.reshape(nb, 1, -1).permute(0, 2, 1) |
| estimates = estimates.reshape(nb, 1, -1).permute(0, 2, 1) |
|
|
| |
| if self.zero_mean: |
| mean_reference = references.mean(dim=1, keepdim=True) |
| mean_estimate = estimates.mean(dim=1, keepdim=True) |
| else: |
| mean_reference = 0 |
| mean_estimate = 0 |
|
|
| _references = references - mean_reference |
| _estimates = estimates - mean_estimate |
|
|
| references_projection = (_references**2).sum(dim=-2) + eps |
| references_on_estimates = (_estimates * _references).sum(dim=-2) + eps |
|
|
| scale = ( |
| (references_on_estimates / references_projection).unsqueeze(1) |
| if self.scaling |
| else 1 |
| ) |
|
|
| e_true = scale * _references |
| e_res = _estimates - e_true |
|
|
| signal = (e_true**2).sum(dim=1) |
| noise = (e_res**2).sum(dim=1) |
| sdr = -10 * torch.log10(signal / noise + eps) |
|
|
| if self.clip_min is not None: |
| sdr = torch.clamp(sdr, min=self.clip_min) |
|
|
| if self.reduction == "mean": |
| sdr = sdr.mean() |
| elif self.reduction == "sum": |
| sdr = sdr.sum() |
| return sdr |
|
|
