linter fixes

This commit is contained in:
Eren Gölge 2021-04-07 12:36:03 +02:00
parent 9782d9ea5d
commit f890454de3
3 changed files with 122 additions and 48 deletions

View File

@ -20,7 +20,7 @@ class Synthesizer(object):
"""General 🐸 TTS interface for inference. It takes a tts and a vocoder
model and synthesize speech from the provided text.
The text is divided into a list of sentences using `pysbd` and synthesize
The text is divided into a list of sentences using `pysbd` and synthesize
speech on each sentence separately.
If you have certain special characters in your text, you need to handle

View File

@ -77,8 +77,8 @@ class GANDataset(Dataset):
"""Pad samples shorter than the output sequence length"""
if len(audio) < self.seq_len:
audio = np.pad(audio, (0, self.seq_len - len(audio)),
mode='constant',
constant_values=0.0)
mode='constant',
constant_values=0.0)
if mel is not None and mel.shape[1] < self.feat_frame_len:
pad_value = self.ap.melspectrogram(np.zeros([self.ap.win_length]))[:, 0]

View File

@ -1,36 +1,67 @@
import torch
import librosa
from torch import nn
from torch.nn import functional as F
class TorchSTFT(nn.Module): # pylint: disable=abstract-method
# TODO: move this to audio.py with a transparent torch API.
def __init__(self, n_fft, hop_length, win_length, pad_mode='reflect', window='hann_window'):
class TorchSTFT(nn.Module):
"""TODO: Merge this with audio.py"""
def __init__(self,
n_fft,
hop_length,
win_length,
window='hann_window',
sample_rate=None,
mel_fmin=0,
mel_fmax=None,
n_mels=80,
use_mel=False):
""" Torch based STFT operation """
super(TorchSTFT, self).__init__()
self.n_fft = n_fft
self.hop_length = hop_length
self.win_length = win_length
self.pad_mode = pad_mode
self.sample_rate = sample_rate
self.mel_fmin = mel_fmin
self.mel_fmax = mel_fmax
self.n_mels = n_mels
self.use_mel = use_mel
self.window = nn.Parameter(getattr(torch, window)(win_length),
requires_grad=False)
self.mel_basis = None
if use_mel:
self._build_mel_basis()
def __call__(self, x):
padding = int((self.n_fft - self.hop_length) / 2)
x = torch.nn.functional.pad(x, (padding, padding), mode='reflect')
# B x D x T x 2
o = torch.stft(x,
self.n_fft,
self.hop_length,
self.win_length,
self.window,
center=True,
pad_mode=self.pad_mode, # needs to be compatible with audio.py
normalized=False,
onesided=True,
return_complex=False)
o = torch.stft(
x.squeeze(1),
self.n_fft,
self.hop_length,
self.win_length,
self.window,
center=True,
pad_mode="reflect", # compatible with audio.py
normalized=False,
onesided=True,
return_complex=False)
M = o[:, :, :, 0]
P = o[:, :, :, 1]
return torch.sqrt(torch.clamp(M ** 2 + P ** 2, min=1e-8))
S = torch.sqrt(torch.clamp(M**2 + P**2, min=1e-8))
if self.use_mel:
S = torch.matmul(self.mel_basis.to(x), S)
return S
def _build_mel_basis(self):
mel_basis = librosa.filters.mel(self.sample_rate,
self.n_fft,
n_mels=self.n_mels,
fmin=self.mel_fmin,
fmax=self.mel_fmax)
self.mel_basis = torch.from_numpy(mel_basis).float()
#################################
@ -39,7 +70,9 @@ class TorchSTFT(nn.Module): # pylint: disable=abstract-method
class STFTLoss(nn.Module):
""" Single scale STFT Loss """
""" STFT loss. Input generate and real waveforms are converted
to spectrograms compared with L1 and Spectral convergence losses.
It is from ParallelWaveGAN paper https://arxiv.org/pdf/1910.11480.pdf"""
def __init__(self, n_fft, hop_length, win_length):
super(STFTLoss, self).__init__()
self.n_fft = n_fft
@ -57,7 +90,9 @@ class STFTLoss(nn.Module):
return loss_mag, loss_sc
class MultiScaleSTFTLoss(torch.nn.Module):
""" Multi scale STFT loss """
""" Multi-scale STFT loss. Input generate and real waveforms are converted
to spectrograms compared with L1 and Spectral convergence losses.
It is from ParallelWaveGAN paper https://arxiv.org/pdf/1910.11480.pdf"""
def __init__(self,
n_ffts=(1024, 2048, 512),
hop_lengths=(120, 240, 50),
@ -79,9 +114,30 @@ class MultiScaleSTFTLoss(torch.nn.Module):
loss_mag /= N
return loss_mag, loss_sc
class L1SpecLoss(nn.Module):
""" L1 Loss over Spectrograms as described in HiFiGAN paper https://arxiv.org/pdf/2010.05646.pdf"""
def __init__(self, sample_rate, n_fft, hop_length, win_length, mel_fmin=None, mel_fmax=None, n_mels=None, use_mel=True):
super().__init__()
self.use_mel = use_mel
self.stft = TorchSTFT(n_fft,
hop_length,
win_length,
sample_rate=sample_rate,
mel_fmin=mel_fmin,
mel_fmax=mel_fmax,
n_mels=n_mels,
use_mel=use_mel)
def forward(self, y_hat, y):
y_hat_M = self.stft(y_hat)
y_M = self.stft(y)
# magnitude loss
loss_mag = F.l1_loss(torch.log(y_M), torch.log(y_hat_M))
return loss_mag
class MultiScaleSubbandSTFTLoss(MultiScaleSTFTLoss):
""" Multiscale STFT loss for multi band model outputs """
""" Multiscale STFT loss for multi band model outputs.
From MultiBand-MelGAN paper https://arxiv.org/abs/2005.05106"""
# pylint: disable=no-self-use
def forward(self, y_hat, y):
y_hat = y_hat.view(-1, 1, y_hat.shape[2])
@ -143,9 +199,12 @@ class MelganFeatureLoss(nn.Module):
# pylint: disable=no-self-use
def forward(self, fake_feats, real_feats):
loss_feats = 0
for fake_feat, real_feat in zip(fake_feats, real_feats):
loss_feats += self.loss_func(fake_feat, real_feat)
loss_feats /= len(fake_feats) + len(real_feats)
num_feats = 0
for idx, _ in enumerate(fake_feats):
for fake_feat, real_feat in zip(fake_feats[idx], real_feats[idx]):
loss_feats += self.loss_func(fake_feat, real_feat)
num_feats += 1
loss_feats = loss_feats / num_feats
return loss_feats
@ -198,24 +257,31 @@ def _apply_D_loss(scores_fake, scores_real, loss_func):
class GeneratorLoss(nn.Module):
"""Generator Loss Wrapper. Based on model configuration it sets a right set of loss functions and computes
losses. It allows to experiment with different combinations of loss functions with different models by just
changing configurations.
Args:
C (AttrDict): model configuration.
"""
def __init__(self, C):
""" Compute Generator Loss values depending on training
configuration """
super(GeneratorLoss, self).__init__()
super().__init__()
assert not(C.use_mse_gan_loss and C.use_hinge_gan_loss),\
" [!] Cannot use HingeGANLoss and MSEGANLoss together."
self.use_stft_loss = C.use_stft_loss
self.use_subband_stft_loss = C.use_subband_stft_loss
self.use_mse_gan_loss = C.use_mse_gan_loss
self.use_hinge_gan_loss = C.use_hinge_gan_loss
self.use_feat_match_loss = C.use_feat_match_loss
self.use_stft_loss = C.use_stft_loss if 'use_stft_loss' in C else False
self.use_subband_stft_loss = C.use_subband_stft_loss if 'use_subband_stft_loss' in C else False
self.use_mse_gan_loss = C.use_mse_gan_loss if 'use_mse_gan_loss' in C else False
self.use_hinge_gan_loss = C.use_hinge_gan_loss if 'use_hinge_gan_loss' in C else False
self.use_feat_match_loss = C.use_feat_match_loss if 'use_feat_match_loss' in C else False
self.use_l1_spec_loss = C.use_l1_spec_loss if 'use_l1_spec_loss' in C else False
self.stft_loss_weight = C.stft_loss_weight
self.subband_stft_loss_weight = C.subband_stft_loss_weight
self.mse_gan_loss_weight = C.mse_G_loss_weight
self.hinge_gan_loss_weight = C.hinge_G_loss_weight
self.feat_match_loss_weight = C.feat_match_loss_weight
self.stft_loss_weight = C.stft_loss_weight if 'stft_loss_weight' in C else 0.0
self.subband_stft_loss_weight = C.subband_stft_loss_weight if 'subband_stft_loss_weight' in C else 0.0
self.mse_gan_loss_weight = C.mse_G_loss_weight if 'mse_G_loss_weight' in C else 0.0
self.hinge_gan_loss_weight = C.hinge_G_loss_weight if 'hinde_G_loss_weight' in C else 0.0
self.feat_match_loss_weight = C.feat_match_loss_weight if 'feat_match_loss_weight' in C else 0.0
self.l1_spec_loss_weight = C.l1_spec_loss_weight if 'l1_spec_loss_weight' in C else 0.0
if C.use_stft_loss:
self.stft_loss = MultiScaleSTFTLoss(**C.stft_loss_params)
@ -227,6 +293,9 @@ class GeneratorLoss(nn.Module):
self.hinge_loss = HingeGLoss()
if C.use_feat_match_loss:
self.feat_match_loss = MelganFeatureLoss()
if C.use_l1_spec_loss:
assert C.audio['sample_rate'] == C.l1_spec_loss_params['sample_rate']
self.l1_spec_loss = L1SpecLoss(**C.l1_spec_loss_params)
def forward(self, y_hat=None, y=None, scores_fake=None, feats_fake=None, feats_real=None, y_hat_sub=None, y_sub=None):
gen_loss = 0
@ -235,35 +304,41 @@ class GeneratorLoss(nn.Module):
# STFT Loss
if self.use_stft_loss:
stft_loss_mg, stft_loss_sc = self.stft_loss(y_hat.squeeze(1), y.squeeze(1))
stft_loss_mg, stft_loss_sc = self.stft_loss(y_hat[:, :, :y.size(2)].squeeze(1), y.squeeze(1))
return_dict['G_stft_loss_mg'] = stft_loss_mg
return_dict['G_stft_loss_sc'] = stft_loss_sc
gen_loss += self.stft_loss_weight * (stft_loss_mg + stft_loss_sc)
gen_loss = gen_loss + self.stft_loss_weight * (stft_loss_mg + stft_loss_sc)
# L1 Spec loss
if self.use_l1_spec_loss:
l1_spec_loss = self.l1_spec_loss(y_hat, y)
return_dict['G_l1_spec_loss'] = l1_spec_loss
gen_loss = gen_loss + self.l1_spec_loss_weight * l1_spec_loss
# subband STFT Loss
if self.use_subband_stft_loss:
subband_stft_loss_mg, subband_stft_loss_sc = self.subband_stft_loss(y_hat_sub, y_sub)
return_dict['G_subband_stft_loss_mg'] = subband_stft_loss_mg
return_dict['G_subband_stft_loss_sc'] = subband_stft_loss_sc
gen_loss += self.subband_stft_loss_weight * (subband_stft_loss_mg + subband_stft_loss_sc)
gen_loss = gen_loss + self.subband_stft_loss_weight * (subband_stft_loss_mg + subband_stft_loss_sc)
# multiscale MSE adversarial loss
if self.use_mse_gan_loss and scores_fake is not None:
mse_fake_loss = _apply_G_adv_loss(scores_fake, self.mse_loss)
return_dict['G_mse_fake_loss'] = mse_fake_loss
adv_loss += self.mse_gan_loss_weight * mse_fake_loss
adv_loss = adv_loss + self.mse_gan_loss_weight * mse_fake_loss
# multiscale Hinge adversarial loss
if self.use_hinge_gan_loss and not scores_fake is not None:
hinge_fake_loss = _apply_G_adv_loss(scores_fake, self.hinge_loss)
return_dict['G_hinge_fake_loss'] = hinge_fake_loss
adv_loss += self.hinge_gan_loss_weight * hinge_fake_loss
adv_loss = adv_loss + self.hinge_gan_loss_weight * hinge_fake_loss
# Feature Matching Loss
if self.use_feat_match_loss and not feats_fake:
if self.use_feat_match_loss and not feats_fake is None:
feat_match_loss = self.feat_match_loss(feats_fake, feats_real)
return_dict['G_feat_match_loss'] = feat_match_loss
adv_loss += self.feat_match_loss_weight * feat_match_loss
adv_loss = adv_loss + self.feat_match_loss_weight * feat_match_loss
return_dict['G_loss'] = gen_loss + adv_loss
return_dict['G_gen_loss'] = gen_loss
return_dict['G_adv_loss'] = adv_loss
@ -271,10 +346,9 @@ class GeneratorLoss(nn.Module):
class DiscriminatorLoss(nn.Module):
""" Compute Discriminator Loss values depending on training
configuration """
"""Like ```GeneratorLoss```"""
def __init__(self, C):
super(DiscriminatorLoss, self).__init__()
super().__init__()
assert not(C.use_mse_gan_loss and C.use_hinge_gan_loss),\
" [!] Cannot use HingeGANLoss and MSEGANLoss together."