STFT loss module.
| 11 | |
| 12 | |
| 13 | class STFTLoss(torch.nn.Module): |
| 14 | """STFT loss module.""" |
| 15 | |
| 16 | def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window", |
| 17 | use_mel_loss=False): |
| 18 | """Initialize STFT loss module.""" |
| 19 | super(STFTLoss, self).__init__() |
| 20 | self.fft_size = fft_size |
| 21 | self.shift_size = shift_size |
| 22 | self.win_length = win_length |
| 23 | self.window = getattr(torch, window)(win_length) |
| 24 | self.spectral_convergenge_loss = SpectralConvergengeLoss() |
| 25 | self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss() |
| 26 | self.use_mel_loss = use_mel_loss |
| 27 | self.mel_basis = None |
| 28 | |
| 29 | def forward(self, x, y): |
| 30 | """Calculate forward propagation. |
| 31 | |
| 32 | Args: |
| 33 | x (Tensor): Predicted signal (B, T). |
| 34 | y (Tensor): Groundtruth signal (B, T). |
| 35 | |
| 36 | Returns: |
| 37 | Tensor: Spectral convergence loss value. |
| 38 | Tensor: Log STFT magnitude loss value. |
| 39 | |
| 40 | """ |
| 41 | x_mag = stft(x, self.fft_size, self.shift_size, self.win_length, self.window) |
| 42 | y_mag = stft(y, self.fft_size, self.shift_size, self.win_length, self.window) |
| 43 | if self.use_mel_loss: |
| 44 | if self.mel_basis is None: |
| 45 | self.mel_basis = torch.from_numpy(librosa.filters.mel(22050, self.fft_size, 80)).cuda().T |
| 46 | x_mag = x_mag @ self.mel_basis |
| 47 | y_mag = y_mag @ self.mel_basis |
| 48 | |
| 49 | sc_loss = self.spectral_convergenge_loss(x_mag, y_mag) |
| 50 | mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag) |
| 51 | |
| 52 | return sc_loss, mag_loss |
| 53 | |
| 54 | |
| 55 | class MultiResolutionSTFTLoss(torch.nn.Module): |