Spaces:
Running on Zero
Running on Zero
File size: 4,902 Bytes
a95f6c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import torch
def get_frequency_weighting(freqs, freq_weighting=None):
if freq_weighting is None:
return torch.ones_like(freqs).to(freqs.device)
elif freq_weighting=="sqrt":
return torch.sqrt(freqs)
elif freq_weighting=="exp":
freqs = torch.exp(freqs)
return freqs - freqs[:, 0, :].unsqueeze(-2)
elif freq_weighting=="log":
return torch.log(1+freqs)
elif freq_weighting=="linear":
return freqs
def get_loss(loss_args, operator=None):
if loss_args.name == "none":
return None
if hasattr(loss_args, "loss_1"): #We have a hybrid of multiple losses
return lambda x, x_hat: torch.sum( torch.stack([ get_loss( getattr(loss_args, key), operator=operator)(x, x_hat) for key in list(loss_args.keys()) ]) )
else:
if "stft" in loss_args.name:
def loss_fn(x, x_hat, freq_weighting=None):
X = operator.apply_stft(x)
X_hat = operator.apply_stft(x_hat)
freqs = torch.linspace(0, 1, X.shape[-2]).to(X.device).unsqueeze(-1).unsqueeze(0).expand(X.shape)+1
freqs = get_frequency_weighting(freqs, freq_weighting=loss_args.get("freq_weighting", None))
X = X * freqs
X_hat = X_hat * freqs
if loss_args.name == "l2_stft_sum":
loss= torch.sum((X-X_hat).abs()**2)
elif loss_args.name == "l2_stft_mag_sum":
loss= torch.sum((X.abs()-X_hat.abs())**2)
elif loss_args.name == "l2_stft_logmag_sum":
loss= torch.sum((torch.log10(X.abs()+1e-8)-torch.log10(X_hat.abs()+1e-8))**2)
elif loss_args.name == "l2_comp_stft_sum":
compression_factor = loss_args.get("compression_factor", None)
assert compression_factor is not None and compression_factor > 0. and compression_factor <= 1., f"Compression factor weird: {compression_factor}"
X_comp = (X.abs()+1e-8)**compression_factor * torch.exp(1j*X.angle())
X_hat_comp= (X_hat.abs()+1e-8)**compression_factor * torch.exp(1j*X_hat.angle())
loss= torch.sum((X_comp - X_hat_comp).abs()**2)
elif loss_args.name == "l2_comp_stft_mean":
compression_factor = loss_args.get("compression_factor", None)
assert compression_factor is not None and compression_factor > 0. and compression_factor <= 1., f"Compression factor weird: {compression_factor}"
X_comp = (X.abs()+1e-8)**compression_factor * torch.exp(1j*X.angle())
X_hat_comp= (X_hat.abs()+1e-8)**compression_factor * torch.exp(1j*X_hat.angle())
loss= torch.mean((X_comp - X_hat_comp).abs()**2)
elif loss_args.name == "l2_comp_stft_summean":
compression_factor = loss_args.get("compression_factor", None)
assert compression_factor is not None and compression_factor > 0. and compression_factor <= 1., f"Compression factor weird: {compression_factor}"
X_comp = (X.abs()+1e-8)**compression_factor * torch.exp(1j*X.angle())
X_hat_comp= (X_hat.abs()+1e-8)**compression_factor * torch.exp(1j*X_hat.angle())
loss= torch.mean(torch.sum((X_comp - X_hat_comp).abs()**2, dim=-2))
elif loss_args.name == "l2_log_stft_sum":
X_comp = torch.log(1+X.abs())* torch.exp(1j*X.angle())
X_hat_comp= torch.log(1+X_hat.abs())* torch.exp(1j*X_hat.angle())
loss= torch.sum((X_comp - X_hat_comp).abs()**2)
else:
raise NotImplementedError(f"rec_loss {loss_args.name} not implemented")
weight=loss_args.get("weight", 1.)
return weight*loss
return lambda x, x_hat: loss_fn(x, x_hat)
else:
if loss_args.name == "l2_sum":
def loss_fn(x, x_hat):
loss= torch.sum((x-x_hat)**2)
weight=loss_args.get("weight", 1.)
return weight*loss
elif loss_args.name == "l2_mean":
def loss_fn(x, x_hat):
loss= torch.mean((x-x_hat)**2)
weight=loss_args.get("weight", 1.)
return weight*loss
else:
raise NotImplementedError(f"rec_loss {loss_args.name} not implemented")
return lambda x, x_hat: loss_fn(x, x_hat)
|