decoder/loss.py

from typing import List, Tuple

import torch
import torchaudio
from torch import nn

from decoder.modules import safe_log

import torch.nn.functional as F


class MelSpecReconstructionLoss(nn.Module):
 """
 L1 distance between the mel-scaled magnitude spectrograms of the ground truth sample and the generated sample
 """

 def __init__(
 self, sample_rate: int = 24000, n_fft: int = 1024, hop_length: int = 256, n_mels: int = 100,
 ):
 super().__init__()
 self.mel_spec = torchaudio.transforms.MelSpectrogram(
 sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, center=True, power=1,
 )

 def forward(self, y_hat, y) -> torch.Tensor:
 """
 Args:
 y_hat (Tensor): Predicted audio waveform.
 y (Tensor): Ground truth audio waveform.

 Returns:
 Tensor: L1 loss between the mel-scaled magnitude spectrograms.
 """
 mel_hat = safe_log(self.mel_spec(y_hat))
 mel = safe_log(self.mel_spec(y))

 loss = torch.nn.functional.l1_loss(mel, mel_hat)

 return loss


class GeneratorLoss(nn.Module):
 """
 Generator Loss module. Calculates the loss for the generator based on discriminator outputs.
 """

 def forward(self, disc_outputs: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
 """
 Args:
 disc_outputs (List[Tensor]): List of discriminator outputs.

 Returns:
 Tuple[Tensor, List[Tensor]]: Tuple containing the total loss and a list of loss values from
 the sub-discriminators
 """
 loss = 0
 gen_losses = []
 for dg in disc_outputs:
 l = torch.mean(torch.clamp(1 - dg, min=0))
 gen_losses.append(l)
 loss += l

 return loss, gen_losses


class DiscriminatorLoss(nn.Module):
 """
 Discriminator Loss module. Calculates the loss for the discriminator based on real and generated outputs.
 """

 def forward(
 self, disc_real_outputs: List[torch.Tensor], disc_generated_outputs: List[torch.Tensor]
 ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
 """
 Args:
 disc_real_outputs (List[Tensor]): List of discriminator outputs for real samples.
 disc_generated_outputs (List[Tensor]): List of discriminator outputs for generated samples.

 Returns:
 Tuple[Tensor, List[Tensor], List[Tensor]]: A tuple containing the total loss, a list of loss values from
 the sub-discriminators for real outputs, and a list of
 loss values for generated outputs.
 """
 loss = 0
 r_losses = []
 g_losses = []
 for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
 r_loss = torch.mean(torch.clamp(1 - dr, min=0))
 g_loss = torch.mean(torch.clamp(1 + dg, min=0))
 loss += r_loss + g_loss
 r_losses.append(r_loss.item())
 g_losses.append(g_loss.item())

 return loss, r_losses, g_losses


class FeatureMatchingLoss(nn.Module):
 """
 Feature Matching Loss module. Calculates the feature matching loss between feature maps of the sub-discriminators.
 """

 def forward(self, fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]) -> torch.Tensor:
 """
 Args:
 fmap_r (List[List[Tensor]]): List of feature maps from real samples.
 fmap_g (List[List[Tensor]]): List of feature maps from generated samples.

 Returns:
 Tensor: The calculated feature matching loss.
 """
 loss = 0
 for dr, dg in zip(fmap_r, fmap_g):
 for rl, gl in zip(dr, dg):
 loss += torch.mean(torch.abs(rl - gl))

 return loss

class DACGANLoss(nn.Module):
 """
 Computes a discriminator loss, given a discriminator on
 generated waveforms/spectrograms compared to ground truth
 waveforms/spectrograms. Computes the loss for both the
 discriminator and the generator in separate functions.
 """

 def __init__(self, discriminator):
 super().__init__()
 self.discriminator = discriminator

 def forward(self, fake, real):
 # d_fake = self.discriminator(fake.audio_data)
 # d_real = self.discriminator(real.audio_data)
 d_fake = self.discriminator(fake)
 d_real = self.discriminator(real)
 return d_fake, d_real

 def discriminator_loss(self, fake, real):
 d_fake, d_real = self.forward(fake.clone().detach(), real)

 loss_d = 0
 for x_fake, x_real in zip(d_fake, d_real):
 loss_d += torch.mean(x_fake[-1] ** 2)
 loss_d += torch.mean((1 - x_real[-1]) ** 2)
 return loss_d

 def generator_loss(self, fake, real):
 d_fake, d_real = self.forward(fake, real)

 loss_g = 0
 for x_fake in d_fake:
 loss_g += torch.mean((1 - x_fake[-1]) ** 2)

 loss_feature = 0

 for i in range(len(d_fake)):
 for j in range(len(d_fake[i]) - 1):
 loss_feature += F.l1_loss(d_fake[i][j], d_real[i][j].detach())
 return loss_g, loss_feature