112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class mean_squared_error(nn.Module):
|
|
|
|
def __init__(self,reduction="mean"):
|
|
super().__init__()
|
|
|
|
self.loss_fun = nn.MSELoss(reduction=reduction)
|
|
self.higher_better = False
|
|
|
|
def forward(self,prediction:torch.Tensor, target: torch.Tensor):
|
|
|
|
if prediction.size() != target.size() or target.ndim < 3:
|
|
raise TypeError(f"""Inputs must be of the same shape (batch_size,channels,samples)
|
|
got {prediction.size()} and {target.size()} instead""")
|
|
|
|
return self.loss_fun(prediction, target)
|
|
|
|
class mean_absolute_error(nn.Module):
|
|
|
|
def __init__(self,reduction="mean"):
|
|
super().__init__()
|
|
|
|
self.loss_fun = nn.L1Loss(reduction=reduction)
|
|
self.higher_better = False
|
|
|
|
def forward(self, prediction:torch.Tensor, target: torch.Tensor):
|
|
|
|
if prediction.size() != target.size() or target.ndim < 3:
|
|
raise TypeError(f"""Inputs must be of the same shape (batch_size,channels,samples)
|
|
got {prediction.size()} and {target.size()} instead""")
|
|
|
|
return self.loss_fun(prediction, target)
|
|
|
|
class Si_SDR(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
reduction:str="mean"
|
|
):
|
|
super().__init__()
|
|
if reduction in ["sum","mean",None]:
|
|
self.reduction = reduction
|
|
else:
|
|
raise TypeError("Invalid reduction, valid options are sum, mean, None")
|
|
self.higher_better = False
|
|
|
|
def forward(self,prediction:torch.Tensor, target:torch.Tensor):
|
|
|
|
if prediction.size() != target.size() or target.ndim < 3:
|
|
raise TypeError(f"""Inputs must be of the same shape (batch_size,channels,samples)
|
|
got {prediction.size()} and {target.size()} instead""")
|
|
|
|
target_energy = torch.sum(target**2,keepdim=True,dim=-1)
|
|
scaling_factor = torch.sum(prediction*target,keepdim=True,dim=-1) / target_energy
|
|
target_projection = target * scaling_factor
|
|
noise = prediction - target_projection
|
|
ratio = torch.sum(target_projection**2,dim=-1) / torch.sum(noise**2,dim=-1)
|
|
si_sdr = 10*torch.log10(ratio).mean(dim=-1)
|
|
|
|
if self.reduction == "sum":
|
|
si_sdr = si_sdr.sum()
|
|
elif self.reduction == "mean":
|
|
si_sdr = si_sdr.mean()
|
|
else:
|
|
pass
|
|
|
|
return si_sdr
|
|
|
|
|
|
|
|
class Avergeloss(nn.Module):
|
|
|
|
def __init__(self,losses):
|
|
super().__init__()
|
|
|
|
self.valid_losses = nn.ModuleList()
|
|
|
|
direction = [getattr(LOSS_MAP[loss](),"higher_better") for loss in losses]
|
|
if len(set(direction)) > 1:
|
|
raise ValueError("all cost functions should be of same nature, maximize or minimize!")
|
|
|
|
self.higher_better = direction[0]
|
|
for loss in losses:
|
|
loss = self.validate_loss(loss)
|
|
self.valid_losses.append(loss())
|
|
|
|
|
|
def validate_loss(self,loss:str):
|
|
if loss not in LOSS_MAP.keys():
|
|
raise ValueError(f"Invalid loss function {loss}, available loss functions are {tuple([loss for loss in LOSS_MAP.keys()])}")
|
|
else:
|
|
return LOSS_MAP[loss]
|
|
|
|
def forward(self,prediction:torch.Tensor, target:torch.Tensor):
|
|
loss = 0.0
|
|
for loss_fun in self.valid_losses:
|
|
loss += loss_fun(prediction, target)
|
|
|
|
return loss
|
|
|
|
|
|
|
|
|
|
LOSS_MAP = {"mae":mean_absolute_error,
|
|
"mse": mean_squared_error,
|
|
"SI-SDR":Si_SDR}
|
|
|
|
|