class flamo.optimize.loss.sparsity_loss(*args, **kwargs)

Calculates the sparsity loss for a given model.

The sparsity loss is calculated based on the feedback loop of the FDN model’s core. It measures the sparsity of the feedback matrix A of size (N, N). Note: for the loss to be compatible with the class flamo.optimize.trainer.Trainer, it requires y_pred and y_target as arguments even if these are not being considered. If the feedback matrix has a third dimension C, A.size = (C, N, N), the loss is calculated as the mean of the contribution of each (N,N) matrix.

\[\mathcal{L} = \frac{\sum_{i,j} |A_{i,j}| - N\sqrt{N}}{N(1 - \sqrt{N})}\]

For more details, refer to the paper Optimizing Tiny Colorless Feedback Delay Networks by Dal Santo, G. et al.

Arguments:

  • y_pred (torch.Tensor): The predicted output.

  • y_target (torch.Tensor): The target output.

  • model (nn.Module): The model containing the core with the feedback loop.

Returns:

torch.Tensor: The calculated sparsity loss.

forward(y_pred: Tensor, y_target: Tensor, model: Module)

Define 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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class flamo.optimize.loss.mse_loss(nfft: int = None, device: str = 'cpu')

Wrapper for the mean squared error loss.

\[\mathcal{L} = \frac{1}{N} \sum_{i=1}^{N} \left( y_{\text{pred},i} - y_{\text{true},i} \right)^2\]

where \(N\) is the number of nfft points and \(M\) is the number of channels.

Arguments / Attributes:

  • nfft (int): Number of FFT points.

  • device (str): Device to run the calculations on.

forward(y_pred, y_true)

Calculates the mean squared error loss. If is_masked is set to True, the loss is calculated using a masked version of the predicted output. This option is useful to introduce stochasticity, as the mask is generated randomly.

Arguments:

  • y_pred (torch.Tensor): The predicted output.

  • y_true (torch.Tensor): The target output.

Returns:

torch.Tensor: The calculated MSE loss.

class flamo.optimize.loss.masked_mse_loss(nfft: int, n_samples: int, n_sets: int = 1, regenerate_mask: bool = True, device: str = 'cpu')

Wrapper for the mean squared error loss with random masking.

Calculates the mean squared error loss between the predicted and target outputs. The loss is calculated using a masked version of the predicted output. This option is useful to introduce stochasticity, as the mask is generated randomly.

\[\mathcal{L} = \frac{1}{\left| \mathbb{S} \right|} \sum_{i \in \mathbb{S}} \left( y_{\text{pred}, i} - y_{\text{true},i} \right)^2\]

where \(\mathbb{S}\) is the set of indices of the mask being analyzed during the training step.

Arguments / Attributes:

  • nfft (int): Number of FFT points.

  • n_samples (int): Number of samples for masking.

  • n_sets (int): Number of sets for masking. Default is 1.

  • regenerate_mask (bool): After all sets are used, if True, the mask is regenerated. Default is True.

  • device (str): Device to run the calculations on. Default is ‘cpu’.

forward(y_pred, y_true)

Calculates the masked mean squared error loss.

Arguments:

  • y_pred (torch.Tensor): The predicted output.

  • y_true (torch.Tensor): The target output.

Returns:

torch.Tensor: The calculated masked MSE loss.