STFT loss module.
| 74 | |
| 75 | |
| 76 | class STFTLoss(torch.nn.Module): |
| 77 | """STFT loss module.""" |
| 78 | |
| 79 | def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window"): |
| 80 | """Initialize STFT loss module.""" |
| 81 | super(STFTLoss, self).__init__() |
| 82 | self.fft_size = fft_size |
| 83 | self.shift_size = shift_size |
| 84 | self.win_length = win_length |
| 85 | self.window = getattr(torch, window)(win_length) |
| 86 | self.spectral_convergenge_loss = SpectralConvergengeLoss() |
| 87 | self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss() |
| 88 | |
| 89 | def forward(self, x, y): |
| 90 | """Calculate forward propagation. |
| 91 | |
| 92 | Args: |
| 93 | x (Tensor): Predicted signal (B, T). |
| 94 | y (Tensor): Groundtruth signal (B, T). |
| 95 | |
| 96 | Returns: |
| 97 | Tensor: Spectral convergence loss value. |
| 98 | Tensor: Log STFT magnitude loss value. |
| 99 | |
| 100 | """ |
| 101 | x_mag = stft(x, self.fft_size, self.shift_size, self.win_length, self.window) |
| 102 | y_mag = stft(y, self.fft_size, self.shift_size, self.win_length, self.window) |
| 103 | sc_loss = self.spectral_convergenge_loss(x_mag, y_mag) |
| 104 | mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag) |
| 105 | |
| 106 | return sc_loss, mag_loss |
| 107 | |
| 108 | |
| 109 | class MultiResolutionSTFTLoss(torch.nn.Module): |