Module audiocraft.losses.sisnr

Classes

class SISNR (sample_rate: int = 16000,
segment: float | None = 20,
overlap: float = 0.5,
epsilon: float = 1.1920928955078125e-07)

Expand source code

class SISNR(nn.Module):
    """SISNR loss.

    Input should be [B, C, T], output is scalar.

    ..Warning:: This function returns the opposite of the SI-SNR (e.g. `-1 * regular_SI_SNR`).
        Consequently, lower scores are better in terms of reconstruction quality,
        in particular, it should be negative if training goes well. This done this way so
        that this module can also be used as a loss function for training model.

    Args:
        sample_rate (int): Sample rate.
        segment (float or None): Evaluate on chunks of that many seconds. If None, evaluate on
            entire audio only.
        overlap (float): Overlap between chunks, i.e. 0.5 = 50 % overlap.
        epsilon (float): Epsilon value for numerical stability.
    """
    def __init__(
        self,
        sample_rate: int = 16000,
        segment: tp.Optional[float] = 20,
        overlap: float = 0.5,
        epsilon: float = torch.finfo(torch.float32).eps,
    ):
        super().__init__()
        self.sample_rate = sample_rate
        self.segment = segment
        self.overlap = overlap
        self.epsilon = epsilon

    def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:
        B, C, T = ref_sig.shape
        assert ref_sig.shape == out_sig.shape

        if self.segment is None:
            frame = T
            stride = T
        else:
            frame = int(self.segment * self.sample_rate)
            stride = int(frame * (1 - self.overlap))

        epsilon = self.epsilon * frame  # make epsilon prop to frame size.

        gt = _unfold(ref_sig, frame, stride)
        est = _unfold(out_sig, frame, stride)
        if self.segment is None:
            assert gt.shape[-1] == 1

        gt = _center(gt)
        est = _center(est)
        dot = torch.einsum("bcft,bcft->bcf", gt, est)

        proj = dot[:, :, :, None] * gt / (epsilon + _norm2(gt))
        noise = est - proj

        sisnr = 10 * (
            torch.log10(epsilon + _norm2(proj)) - torch.log10(epsilon + _norm2(noise))
        )
        return -1 * sisnr[..., 0].mean()

SISNR loss.

Input should be [B, C, T], output is scalar.

Warning: This function returns the opposite of the SI-SNR (e.g. -1 * regular_SI_SNR).

Consequently, lower scores are better in terms of reconstruction quality, in particular, it should be negative if training goes well. This done this way so that this module can also be used as a loss function for training model.

Args

sample_rate : int

Sample rate.

segment : float or None

Evaluate on chunks of that many seconds. If None, evaluate on entire audio only.

overlap : float

Overlap between chunks, i.e. 0.5 = 50 % overlap.

epsilon : float

Epsilon value for numerical stability.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Class variables

var call_super_init : bool

var dump_patches : bool

var training : bool

Methods

def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) ‑> torch.Tensor

Expand source code

def forward(self, out_sig: torch.Tensor, ref_sig: torch.Tensor) -> torch.Tensor:
    B, C, T = ref_sig.shape
    assert ref_sig.shape == out_sig.shape

    if self.segment is None:
        frame = T
        stride = T
    else:
        frame = int(self.segment * self.sample_rate)
        stride = int(frame * (1 - self.overlap))

    epsilon = self.epsilon * frame  # make epsilon prop to frame size.

    gt = _unfold(ref_sig, frame, stride)
    est = _unfold(out_sig, frame, stride)
    if self.segment is None:
        assert gt.shape[-1] == 1

    gt = _center(gt)
    est = _center(est)
    dot = torch.einsum("bcft,bcft->bcf", gt, est)

    proj = dot[:, :, :, None] * gt / (epsilon + _norm2(gt))
    noise = est - proj

    sisnr = 10 * (
        torch.log10(epsilon + _norm2(proj)) - torch.log10(epsilon + _norm2(noise))
    )
    return -1 * sisnr[..., 0].mean()

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.