From f45b214593b0c4bf4745c4e56daebec21616618a Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Fri, 4 Oct 2024 16:18:45 -0400 Subject: [PATCH] added multiscale mel --- custom_nodes/rvc_nodes.py | 2 + lib/train/losses.py | 247 +++++++++++++++++++++++++++++++++----- training_cli.py | 23 ++-- 3 files changed, 232 insertions(+), 40 deletions(-) diff --git a/custom_nodes/rvc_nodes.py b/custom_nodes/rvc_nodes.py index 15b62b1..dfb467a 100644 --- a/custom_nodes/rvc_nodes.py +++ b/custom_nodes/rvc_nodes.py @@ -359,6 +359,7 @@ class RVCTrainParamsNode: return { "optional": dict( batch_size=("INT",dict(default=4,min=1,max=64,step=1)), + c_gen=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)), c_mel=("FLOAT",dict(default=45.,min=0.,max=100.,step=.1)), c_kl=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)), c_fm=("FLOAT",dict(default=2.,min=0.,max=100.,step=.1)), @@ -366,6 +367,7 @@ class RVCTrainParamsNode: c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), c_tsi=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), c_gp=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)), + use_multiscale=("BOOLEAN",dict(default=False)), use_balancer=("BOOLEAN",dict(default=False)), use_pareto=("BOOLEAN",dict(default=False)), fast_mode=("BOOLEAN",dict(default=False)), diff --git a/lib/train/losses.py b/lib/train/losses.py index 501458c..9392650 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -1,4 +1,10 @@ -from typing import Optional +from collections import namedtuple +import functools +from typing import Callable, List, Optional +import typing +from librosa.filters import mel as librosa_mel_fn +import numpy as np +from scipy import signal import torch import torchaudio.transforms as T import torch.nn.functional as F @@ -43,7 +49,7 @@ class LossBalancer: use_norm = self.use_norm ) - def update_ema_weights(self, new_weights): + def update_ema_weights(self, new_weights: dict): """Updates the EMA of the weights.""" if not self.ema_weights: self.ema_weights = dict(**new_weights) else: self.ema_weights = { @@ -56,19 +62,21 @@ class LossBalancer: """Updates the EMA of the weights.""" if not self.historical_losses: self.historical_losses = dict(**new_losses) else: - for k in new_losses: - self.historical_losses[k]=self.loss_decay*self.historical_losses.get(k, 0.) + (1 - self.loss_decay)*new_losses[k] + for k,v in new_losses.items(): + if np.nan_to_num(v,nan=-1)==-1: print(f"{k=} {v=}") + self.historical_losses[k]=self.loss_decay*self.historical_losses.get(k, v) + (1 - self.loss_decay)*v return dict(**self.historical_losses) - def calculate_loss_slope(self, key, current_loss): + def calculate_loss_slope(self, key: str, current_loss: torch.Tensor): """Calculates the slope of the loss using the current loss and its EMA.""" ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon slope = (current_loss-ema_loss)/ema_loss # relative loss change return slope - def calculate_gradients(self, key, current_loss, input): + def calculate_gradients(self, key: str, current_loss: torch.Tensor, input: torch.Tensor): """Calculates the gradient norm of the current loss wrt the model params.""" self.model.zero_grad() # Clear previous gradients + input.requires_grad_(True) # Backward pass to output layer output_params = torch.autograd.grad(current_loss, [input], retain_graph=True, allow_unused=True, materialize_grads=True)[0] @@ -89,7 +97,7 @@ class LossBalancer: top_20_count = max(int(0.2 * len(sorted_data)),1) # Calculate the weights - weights = [0] * len(sorted_data) + weights = np.zeros(len(sorted_data)) total_weight = 1.0 top_20_weight = 0.8 * total_weight remaining_weight = total_weight - top_20_weight @@ -141,14 +149,13 @@ class LossBalancer: if self.use_norm and input is not None: # Compute L2 gradient norm - input.requires_grad_(True) grad_norm = self.calculate_gradients(key, weighted_loss, input) gradients[key] = max(grad_norm.item(), self.epsilon) else: # Use loss plateau detection with EMA loss_slope = self.calculate_loss_slope(key, weighted_loss) gradients[key] = max(loss_slope.item(), self.epsilon) - valid_losses[key] = weighted_loss + valid_losses[key] = weighted_loss.nan_to_num(self.epsilon) if not valid_losses: return torch.tensor(0.0) # If all losses are skipped @@ -163,10 +170,10 @@ class LossBalancer: self.update_ema_weights(normalized_weights) balanced_loss = 0 for k, loss in valid_losses.items(): - balanced_loss += self.ema_weights[k] * loss + balanced_loss += self.ema_weights.get(k, 1.) * loss - # update historical losses - self.update_historical_losses({k: loss.item()}) + # update historical losses + self.update_historical_losses({k: loss.item() for k, loss in valid_losses.items()}) return balanced_loss @@ -185,7 +192,7 @@ class LossBalancer: @property def weighted_ema_loss(self): - return sum(v*self.ema_weights.get(k,1./len(self.historical_losses)) for k,v in self.historical_losses.items()) + return sum(v*self.ema_weights.get(k,1.) for k,v in self.historical_losses.items()) def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitude: torch.Tensor, dim=-1, eps=1e-8): """ @@ -229,7 +236,7 @@ def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3, eps=1e- max_filtered_tensor = F.max_pool1d(log_magnitude, kernel_size=kernel_size, stride=1, padding=kernel_size//2) # Replace NaNs with zeros and sum along the second-to-last dimension - return max_filtered_tensor.nan_to_num(0).sum(dim) + return max_filtered_tensor.nan_to_num(eps).sum(dim) def compute_tefs(audio_signal: torch.Tensor, eps=1e-8): """ @@ -274,7 +281,7 @@ def compute_tefs(audio_signal: torch.Tensor, eps=1e-8): # Compute the instantaneous phase from the analytic signal phase = torch.angle(analytic_signal).cos() - return envelope.nan_to_num(0), phase.nan_to_num(0) + return envelope.nan_to_num(eps), phase.nan_to_num(eps) def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-8): # Calculate log-magnitude spectrograms @@ -283,10 +290,10 @@ def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percu percussive_list = [] for kernel_size in harmonic_kernel_sizes: - harmonic_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-1).nan_to_num(0).view(mag.size(0),-1)) + harmonic_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-1).nan_to_num(eps).view(mag.size(0),-1)) for kernel_size in percussive_kernel_sizes: - percussive_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-2).nan_to_num(0).view(mag.size(0),-1)) + percussive_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-2).nan_to_num(eps).view(mag.size(0),-1)) # Concatenate the results harmonic = torch.cat(harmonic_list, dim=-1) @@ -296,7 +303,7 @@ def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percu harmonic = minmax_scale(harmonic, eps=eps) percussive = minmax_scale(percussive, eps=eps) - return harmonic.nan_to_num(0), percussive.nan_to_num(0) + return harmonic.nan_to_num(eps), percussive.nan_to_num(eps) def combined_aux_loss( original_audio: torch.Tensor, generated_audio: torch.Tensor, @@ -392,41 +399,223 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens loss += torch.log1p((grad_norm - 1) ** 2) return loss/len(disc_interpolated_output) -def feature_loss(fmap_r, fmap_g): +# Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/loss.py +# LICENSE: https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE +class MultiScaleMelSpectrogramLoss(torch.nn.Module): + """Compute distance between mel spectrograms. Can be used + in a multi-scale way. + + Parameters + ---------- + n_mels : List[int] + Number of mels per STFT, by default [5, 10, 20, 40, 80, 160, 320], + window_lengths : List[int], optional + Length of each window of each STFT, by default [32, 64, 128, 256, 512, 1024, 2048] + loss_fn : typing.Callable, optional + How to compare each loss, by default nn.L1Loss() + clamp_eps : float, optional + Clamp on the log magnitude, below, by default 1e-5 + mag_weight : float, optional + Weight of raw magnitude portion of loss, by default 0.0 (no ampliciation on mag part) + log_weight : float, optional + Weight of log magnitude portion of loss, by default 1.0 + pow : float, optional + Power to raise magnitude to before taking log, by default 1.0 + weight : float, optional + Weight of this loss, by default 1.0 + match_stride : bool, optional + Whether to match the stride of convolutional layers, by default False + + Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py + Additional code copied and modified from https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py + """ + + def __init__( + self, + sampling_rate: int, + n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320], + window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048], + loss_fn: Callable = torch.nn.L1Loss(), + clamp_eps: float = 1e-5, + mag_weight: float = 0.0, + log_weight: float = 1.0, + pow: float = 1.0, + weight: float = 1.0, + match_stride: bool = False, + mel_fmin: List[float] = [0, 0, 0, 0, 0, 0, 0], + mel_fmax: List[float] = [None, None, None, None, None, None, None], + window_type: str = "hann", + ): + super().__init__() + self.sampling_rate = sampling_rate + + STFTParams = namedtuple( + "STFTParams", + ["window_length", "hop_length", "window_type", "match_stride"], + ) + + self.stft_params = [ + STFTParams( + window_length=w, + hop_length=w // 4, + match_stride=match_stride, + window_type=window_type, + ) + for w in window_lengths + ] + self.n_mels = n_mels + self.loss_fn = loss_fn + self.clamp_eps = clamp_eps + self.log_weight = log_weight + self.mag_weight = mag_weight + self.weight = weight + self.mel_fmin = mel_fmin + self.mel_fmax = mel_fmax + self.pow = pow + + @staticmethod + @functools.lru_cache(None) + def get_window( + window_type, + window_length, + ): + return signal.get_window(window_type, window_length) + + @staticmethod + @functools.lru_cache(None) + def get_mel_filters(sr, n_fft, n_mels, fmin, fmax): + return librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax) + + def mel_spectrogram( + self, + wav, + n_mels, + fmin, + fmax, + window_length, + hop_length, + match_stride, + window_type, + ): + """ + Mirrors AudioSignal.mel_spectrogram used by BigVGAN-v2 training from: + https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py + """ + B, C, T = wav.shape + + if match_stride: + assert ( + hop_length == window_length // 4 + ), "For match_stride, hop must equal n_fft // 4" + right_pad = np.ceil(T / hop_length) * hop_length - T + pad = (window_length - hop_length) // 2 + else: + right_pad = 0 + pad = 0 + + wav = torch.nn.functional.pad(wav, (pad, pad + right_pad), mode="reflect") + + window = self.get_window(window_type, window_length) + window = torch.from_numpy(window).to(wav.device).float() + + stft = torch.stft( + wav.reshape(-1, T), + n_fft=window_length, + hop_length=hop_length, + window=window, + return_complex=True, + center=True, + ) + _, nf, nt = stft.shape + stft = stft.reshape(B, C, nf, nt) + if match_stride: + """ + Drop first two and last two frames, which are added, because of padding. Now num_frames * hop_length = num_samples. + """ + stft = stft[..., 2:-2] + magnitude = torch.abs(stft) + + nf = magnitude.shape[2] + mel_basis = self.get_mel_filters( + self.sampling_rate, 2 * (nf - 1), n_mels, fmin, fmax + ) + mel_basis = torch.from_numpy(mel_basis).to(wav.device) + mel_spectrogram = magnitude.transpose(2, -1) @ mel_basis.T + mel_spectrogram = mel_spectrogram.transpose(-1, 2) + + return mel_spectrogram + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Computes mel loss between an estimate and a reference + signal. + + Parameters + ---------- + x : torch.Tensor + Estimate signal + y : torch.Tensor + Reference signal + + Returns + ------- + torch.Tensor + Mel loss. + """ + + loss = 0.0 + for n_mels, fmin, fmax, s in zip( + self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params + ): + kwargs = { + "n_mels": n_mels, + "fmin": fmin, + "fmax": fmax, + "window_length": s.window_length, + "hop_length": s.hop_length, + "match_stride": s.match_stride, + "window_type": s.window_type, + } + + x_mels = self.mel_spectrogram(x, **kwargs) + y_mels = self.mel_spectrogram(y, **kwargs) + x_logmels = torch.log10(x_mels.pow(self.pow)+self.clamp_eps).nan_to_num(self.clamp_eps) + y_logmels = torch.log10(y_mels.pow(self.pow)+self.clamp_eps).nan_to_num(self.clamp_eps) + + if self.log_weight!=0: loss += self.log_weight * self.loss_fn(x_logmels, y_logmels) + if self.mag_weight!=0: loss += self.mag_weight * self.loss_fn(x_mels, y_mels) + + return loss + +def feature_loss(fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]): loss = 0 for dr, dg in zip(fmap_r, fmap_g): for rl, gl in zip(dr, dg): - rl = rl.float() - gl = gl.float() loss += torch.mean(torch.abs(rl - gl)) - return loss -def discriminator_loss(disc_real_outputs, disc_generated_outputs): +def discriminator_loss( + disc_real_outputs: List[torch.Tensor], + disc_generated_outputs: List[torch.Tensor] + ): loss = 0 r_losses = [] g_losses = [] for dr, dg in zip(disc_real_outputs, disc_generated_outputs): - dr = dr.float() - dg = dg.float() r_loss = torch.mean((1 - dr) ** 2) g_loss = torch.mean(dg**2) loss += r_loss + g_loss r_losses.append(r_loss.item()) g_losses.append(g_loss.item()) - # loss /= len(disc_generated_outputs) #average return loss, r_losses, g_losses -def generator_loss(disc_outputs): +def generator_loss(disc_outputs: List[torch.Tensor]): loss = 0 gen_losses = [] for dg in disc_outputs: - dg = dg.float() l = torch.mean((1 - dg) ** 2) gen_losses.append(l) loss += l - # loss /= len(disc_outputs) #average return loss, gen_losses diff --git a/training_cli.py b/training_cli.py index 29d383d..77b7157 100644 --- a/training_cli.py +++ b/training_cli.py @@ -31,7 +31,7 @@ from .lib.train.data_utils import ( TextAudioCollate, DistributedBucketSampler, ) -from .lib.train.losses import LossBalancer, combined_aux_loss, generator_loss, discriminator_loss, feature_loss, gradient_norm_loss, kl_loss +from .lib.train.losses import LossBalancer, MultiScaleMelSpectrogramLoss, combined_aux_loss, generator_loss, discriminator_loss, feature_loss, gradient_norm_loss, kl_loss from .lib.train.mel_processing import mel_spectrogram_torch, spec_to_mel_torch def save_checkpoint(ckpt, name, epoch, hps, model_path=None): @@ -101,7 +101,6 @@ def train_model(hps: "utils.HParams"): n_gpus = 1 gpu_devices = hps.gpus.split("-") if hps.gpus else range(n_gpus) - children = {} for i, device in enumerate(gpu_devices): @@ -122,7 +121,8 @@ def train_model(hps: "utils.HParams"): def run(rank, n_gpus, hps, device): print(f"{__name__=}") - global global_step, least_loss, loss_file, best_model_name, gradient_clip_value + global global_step, least_loss, loss_file, best_model_name, gradient_clip_value, MultiscaleMelLoss + MultiscaleMelLoss = MultiScaleMelSpectrogramLoss(hps.data.sampling_rate) global_step = 0 loss_file = os.path.join(hps.model_dir,"losses.json") gradient_clip_value = np.log(hps.train.c_gp/hps.train.learning_rate) if hps.train.get("c_gp",0.)>0 else None @@ -297,7 +297,7 @@ def run(rank, n_gpus, hps, device): use_pareto=hps.train.get("use_pareto",False), use_norm=not hps.train.get("fast_mode",False), initial_weights=dict( - loss_gen=1., + loss_gen=hps.train.get("c_gen",1.), loss_fm=hps.train.get("c_fm",2.), loss_mel=hps.train.get("c_mel",45.), loss_kl=hps.train.get("c_kl",1.), @@ -351,7 +351,7 @@ def train_and_evaluate( if writers is not None: writer, _ = writers - global global_step, least_loss, loss_file, best_model_name, gradient_clip_value + global global_step, least_loss, loss_file, best_model_name, gradient_clip_value, MultiscaleMelLoss net_g.train() net_d.train() @@ -532,9 +532,10 @@ def train_and_evaluate( # Generator y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat) with autocast(enabled=False): - loss_mel = F.l1_loss(y_mel, y_hat_mel) #*hps.train.get("c_mel",0.) - loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) #*hps.train.get("c_kl",0.) - loss_fm = feature_loss(fmap_r, fmap_g) #*hps.train.get("c_fm",0.) + if hps.train.get("use_multiscale"): loss_mel = MultiscaleMelLoss(y_hat, wave) + else: loss_mel = F.l1_loss(y_mel, y_hat_mel) + loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) + loss_fm = feature_loss(fmap_r, fmap_g) harmonic_loss, tefs_loss, tsi_loss = combined_aux_loss( wave, y_hat,n_mels=hps.data.n_mel_channels,sample_rate=hps.data.sampling_rate, c_tefs=hps.train.get("c_tefs",0.), @@ -595,7 +596,7 @@ def train_and_evaluate( **{f"loss/g/{i}": v for i, v in enumerate(losses_gen)}, **{f"loss/d_r/{i}": v for i, v in enumerate(losses_disc_r)}, **{f"loss/d_g/{i}": v for i, v in enumerate(losses_disc_g)}, - **{f"balancer/weights/{k}": v for k, v in balancer.ema_weights.items()} + **{f"balancer/weights/{k}": v for k, v in balancer.ema_weights.items()}, **{f"balancer/losses/{k}": v for k, v in balancer.historical_losses.items()} } @@ -603,7 +604,7 @@ def train_and_evaluate( "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()), "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()), - "slice/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).abs().data.cpu().numpy(), cmap="hot") + "slice/diff^2": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).pow(2).data.cpu().numpy(), cmap="hot") } with torch.no_grad(): @@ -643,7 +644,7 @@ def train_and_evaluate( optim_d, hps.train.learning_rate, epoch, - os.path.join(hps.model_dir, f"D_{saved_epoch}.pth"), + os.path.join(hps.model_dir, f"D_{saved_epoch}.pth") ) if hps.save_every_weights: ckpt = net_g.module.state_dict() if hasattr(net_g, "module") else net_g.state_dict()