From a76cbc1803e52414dcab30d562c64fef4e63a9ba Mon Sep 17 00:00:00 2001 From: jstzwj <1103870790@qq.com> Date: Tue, 22 Nov 2022 19:57:05 +0800 Subject: [PATCH] v0.3.0 --- README.md | 2 +- hifigan/hub/__init__.py | 2 +- hifigan/model/hifigan.py | 9 ++++----- hifigan/model/pipeline.py | 10 ++++------ train.py | 9 ++++++--- 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 54ade1d..b6e661e 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ class AudioPipeline(torch.nn.Module): device = "cpu" -hifigan = torch.hub.load("vtuber-plan/hifi-gan:v0.2.1", "hifigan_48k", force_reload=True).to(device) +hifigan = torch.hub.load("vtuber-plan/hifi-gan:v0.3.0", "hifigan_48k", force_reload=True).to(device) # Load audio wav, sr = torchaudio.load("test.wav") diff --git a/hifigan/hub/__init__.py b/hifigan/hub/__init__.py index 6d070d5..b2400e2 100644 --- a/hifigan/hub/__init__.py +++ b/hifigan/hub/__init__.py @@ -1,5 +1,5 @@ CKPT_URLS = { - "hifigan-48k": "https://github.com/vtuber-plan/hifi-gan/releases/download/v0.2.1/hifigan-48k-C8FDBD55FE7700384955A6EC41AF1D84.pt", + "hifigan-48k": "https://github.com/vtuber-plan/hifi-gan/releases/download/v0.3.0/hifigan-48k-B67C217083569F978E07EFD1AD7B1766.pt", } import torch from ..model.generators.generator import Generator diff --git a/hifigan/model/hifigan.py b/hifigan/model/hifigan.py index f2c27e8..64208d1 100644 --- a/hifigan/model/hifigan.py +++ b/hifigan/model/hifigan.py @@ -48,11 +48,10 @@ def __init__(self, **kwargs): n_fft=self.hparams.data.filter_length, n_mel=self.hparams.data.n_mel_channels, win_length=self.hparams.data.win_length, - hop_length=self.hparams.data.hop_length, - aug=True) + hop_length=self.hparams.data.hop_length) for param in self.audio_pipeline.parameters(): param.requires_grad = False - + # metrics self.valid_mel_loss = torchmetrics.MeanMetric() @@ -61,7 +60,7 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, optimize y_wav, y_wav_lengths = batch["y_wav_values"], batch["y_wav_lengths"] with torch.inference_mode(): - x_mel = self.audio_pipeline(x_wav.squeeze(1)) + x_mel = self.audio_pipeline(x_wav.squeeze(1), aug=True) x_mel_lengths = (x_wav_lengths / self.hparams.data.hop_length).long() x_mel, ids_slice = rand_slice_segments(x_mel, x_mel_lengths, self.hparams.train.segment_size // self.hparams.data.hop_length) @@ -175,7 +174,7 @@ def validation_step(self, batch, batch_idx): y_wav, y_wav_lengths = batch["y_wav_values"], batch["y_wav_lengths"] with torch.inference_mode(): - x_mel = self.audio_pipeline(x_wav.squeeze(1)) + x_mel = self.audio_pipeline(x_wav.squeeze(1), aug=False) x_mel_lengths = (x_wav_lengths / self.hparams.data.hop_length).long() y_spec = spectrogram_torch_audio(y_wav.squeeze(1), diff --git a/hifigan/model/pipeline.py b/hifigan/model/pipeline.py index 1ff0efe..ac2377e 100644 --- a/hifigan/model/pipeline.py +++ b/hifigan/model/pipeline.py @@ -17,13 +17,11 @@ def __init__( n_fft=1024, n_mel=128, win_length=1024, - hop_length=256, - aug=False + hop_length=256 ): super().__init__() self.freq=freq - self.aug=aug pad = int((n_fft-hop_length)/2) self.spec = T.Spectrogram(n_fft=n_fft, win_length=win_length, hop_length=hop_length, @@ -32,18 +30,18 @@ def __init__( # self.strech = T.TimeStretch(hop_length=hop_length, n_freq=freq) self.spec_aug = torch.nn.Sequential( T.FrequencyMasking(freq_mask_param=80), - T.TimeMasking(time_mask_param=80), + # T.TimeMasking(time_mask_param=80), ) self.mel_scale = T.MelScale(n_mels=n_mel, sample_rate=freq, n_stft=n_fft // 2 + 1) - def forward(self, waveform: torch.Tensor) -> torch.Tensor: + def forward(self, waveform: torch.Tensor, aug: bool=False) -> torch.Tensor: shift_waveform = waveform # Convert to power spectrogram spec = self.spec(shift_waveform) spec = torch.sqrt(spec.real.pow(2) + spec.imag.pow(2) + 1e-6) # Apply SpecAugment - if self.aug: + if aug: spec = self.spec_aug(spec) # Convert to mel-scale mel = self.mel_scale(spec) diff --git a/train.py b/train.py index c7cb8ea..024ca2e 100644 --- a/train.py +++ b/train.py @@ -54,8 +54,11 @@ def main(): devices = [int(n.strip()) for n in args.device.split(",")] - checkpoint_callback = ModelCheckpoint(dirpath=None, save_last=True, every_n_train_steps=2000, save_weights_only=False) - earlystop_callback = EarlyStopping(monitor="valid/loss_mel_epoch", mode="min", patience=7) + checkpoint_callback = ModelCheckpoint( + dirpath=None, save_last=True, every_n_train_steps=2000, save_weights_only=False, + monitor="valid/loss_mel_epoch", mode="min", save_top_k=3 + ) + earlystop_callback = EarlyStopping(monitor="valid/loss_mel_epoch", mode="min", patience=3) trainer_params = { "accelerator": args.accelerator, @@ -84,7 +87,7 @@ def main(): batch_per_gpu = hparams.train.batch_size // len(devices) else: batch_per_gpu = hparams.train.batch_size - train_loader = DataLoader(train_dataset, batch_size=batch_per_gpu, num_workers=4, shuffle=True, pin_memory=True, collate_fn=collate_fn) + train_loader = DataLoader(train_dataset, batch_size=batch_per_gpu, num_workers=8, shuffle=True, pin_memory=True, collate_fn=collate_fn) valid_loader = DataLoader(valid_dataset, batch_size=4, num_workers=4, shuffle=False, pin_memory=True, collate_fn=collate_fn) # model