File size: 1,834 Bytes
85a7d2c |
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 |
import numpy as np
import torch
import torch.nn as nn
import torchaudio
from torch.nn import functional as F
from .core import upsample
class SSSLoss(nn.Module):
"""
Single-scale Spectral Loss.
"""
def __init__(self, n_fft=111, alpha=1.0, overlap=0, eps=1e-7):
super().__init__()
self.n_fft = n_fft
self.alpha = alpha
self.eps = eps
self.hop_length = int(n_fft * (1 - overlap)) # 25% of the length
self.spec = torchaudio.transforms.Spectrogram(n_fft=self.n_fft, hop_length=self.hop_length, power=1, normalized=True, center=False)
def forward(self, x_true, x_pred):
S_true = self.spec(x_true) + self.eps
S_pred = self.spec(x_pred) + self.eps
converge_term = torch.mean(torch.linalg.norm(S_true - S_pred, dim = (1, 2)) / torch.linalg.norm(S_true + S_pred, dim = (1, 2)))
log_term = F.l1_loss(S_true.log(), S_pred.log())
loss = converge_term + self.alpha * log_term
return loss
class RSSLoss(nn.Module):
'''
Random-scale Spectral Loss.
'''
def __init__(self, fft_min, fft_max, n_scale, alpha=1.0, overlap=0, eps=1e-7, device='cuda'):
super().__init__()
self.fft_min = fft_min
self.fft_max = fft_max
self.n_scale = n_scale
self.lossdict = {}
for n_fft in range(fft_min, fft_max):
self.lossdict[n_fft] = SSSLoss(n_fft, alpha, overlap, eps).to(device)
def forward(self, x_pred, x_true):
value = 0.
n_ffts = torch.randint(self.fft_min, self.fft_max, (self.n_scale,))
for n_fft in n_ffts:
loss_func = self.lossdict[int(n_fft)]
value += loss_func(x_true, x_pred)
return value / self.n_scale
|