From eb758b5fe099571ee04a278a0964a65776c67448 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Mon, 7 Oct 2024 17:52:57 -0400 Subject: [PATCH] fixed gradient penalty + tweaked balancer --- custom_nodes/audio_nodes.py | 4 +- lib/infer_pack/models.py | 2 - lib/train/losses.py | 77 ++++++++++++++++++++++--------------- training_cli.py | 4 +- 4 files changed, 49 insertions(+), 38 deletions(-) diff --git a/custom_nodes/audio_nodes.py b/custom_nodes/audio_nodes.py index a5ad7f4..8024de2 100644 --- a/custom_nodes/audio_nodes.py +++ b/custom_nodes/audio_nodes.py @@ -233,9 +233,9 @@ class ProcessAudioNode: "optional": { "audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),{"default": None}), "dynamic_threshold_sample_size": ("INT",{"default": 4000, "min": 160, "max": 48000, "step": 160}), - "dynamic_threshold_multiplier": ("FLOAT",{"default": 2.}), + "dynamic_threshold_multiplier": ("FLOAT",{"default": 2., "min": 1.5, "step": .1}), "dynamic_threshold_fill_method": (["median","interpolation"],{"default": "median"}), - "dynamic_threshold_kernel_size": ("INT",{"default": 5}), + "dynamic_threshold_kernel_size": ("INT",{"default": 5, "min": 3, "step": 2}), "silence_threshold_db": ("INT",{"default": -50, "min": -120, "max": 0}), "normalize_threshold_db": ("INT",{"default": -1, "min": -10, "max": 0}) } diff --git a/lib/infer_pack/models.py b/lib/infer_pack/models.py index 9d7df37..574f567 100644 --- a/lib/infer_pack/models.py +++ b/lib/infer_pack/models.py @@ -1041,8 +1041,6 @@ class MultiPeriodDiscriminator(torch.nn.Module): for i, d in enumerate(self.discriminators): y_d_r, fmap_r = d(y) y_d_g, fmap_g = d(y_hat) - # for j in range(len(fmap_r)): - # print(i,j,y.shape,y_hat.shape,fmap_r[j].shape,fmap_g[j].shape) y_d_rs.append(y_d_r) y_d_gs.append(y_d_g) fmap_rs.append(fmap_r) diff --git a/lib/train/losses.py b/lib/train/losses.py index 26f4d37..1c3e738 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -133,6 +133,30 @@ class LossBalancer: normalized_loss_dict = {keys[i]: normalized_losses[i] for i in range(len(keys))} return normalized_loss_dict + + def redistribute_weights(self, gradients: dict): + if self.use_pareto: pareto_weights = self.pareto_normalizer(self.historical_losses) + else: pareto_weights = {} + + inverse_total_gradient = 1./(sum(gradients.values()) + self.epsilon) + total_initial_weight = sum(self.initial_weights.values())-len(gradients) + if total_initial_weight<0: return {k: 1. for k in gradients} + + normalized_weights = {} + + # First pass: sum of Pareto contributions (if any) and weight ratios + for k, g in gradients.items(): + + # Compute the weight ratio based on the gradient + w_ratio = g * inverse_total_gradient + + # average weights + smoothed_ratio = pareto_weights.get(k, w_ratio)*.5 + w_ratio*.5 + + # Store the smoothed ratio for further calculations + normalized_weights[k] = 1. + total_initial_weight * (smoothed_ratio) + + return normalized_weights def on_train_batch_start(self, losses: dict, input: Optional[torch.Tensor]=None): """ @@ -170,23 +194,16 @@ class LossBalancer: # 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.nan_to_num(self.epsilon) + valid_losses[key] = loss.nan_to_num(self.epsilon) if not valid_losses or not gradients: return torch.tensor(0.0) # If all losses are skipped # update historical losses - historical_losses = self.update_historical_losses({k: loss.item() for k, loss in valid_losses.items()}) + self.update_historical_losses({k: loss.item() for k, loss in valid_losses.items()}) # Calculate loss weights based on gradient magnitudes - if len(valid_losses)>1: - if self.use_pareto: pareto_weights = self.pareto_normalizer(historical_losses) - else: pareto_weights = {} - inverse_total_gradient = 1./(sum(gradients.values()) + self.epsilon) - normalized_weights = {} - for k, w in gradients.items(): - w_ratio = w * inverse_total_gradient - normalized_weights[k] = 1. + (w_ratio + pareto_weights.get(k, w_ratio)) / 2. - else: normalized_weights = {k: 1. for k in valid_losses.keys()} + if len(valid_losses)>1: normalized_weights = self.redistribute_weights(gradients) + else: normalized_weights = {k: self.initial_weights.get(k,1.) for k in valid_losses.keys()} # Update EMA weights normalized_weights = self.update_ema_weights(normalized_weights) @@ -382,7 +399,6 @@ def combined_aux_loss( return harmonic_loss, tefs_loss, tsi_loss def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module, eps=1e-8): - loss=0 # Compute the gradient penalty # Randomly interpolate between real and generated data size = [1]*original_audio.ndim @@ -391,22 +407,23 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens interpolated = alpha * original_audio + (1 - alpha) * generated_audio interpolated.requires_grad_(True) # Get the discriminator output for the interpolated data - _, disc_interpolated_output, _, _ = net_d(original_audio, interpolated) + y_d_hat_r, y_d_hat_g, _, _ = net_d(original_audio, interpolated) + loss_disc, _ = discriminator_loss(y_d_hat_r, y_d_hat_g) # Compute gradients of discriminator output w.r.t. interpolated data - for output in disc_interpolated_output: - net_d.zero_grad() - gradients = torch.autograd.grad( - outputs=output, - inputs=interpolated, - grad_outputs=torch.ones_like(output, device=original_audio.device), - retain_graph=True, - allow_unused=True, - materialize_grads=True - )[0] - if gradients.ndim>1: grad_norm = gradients.view(gradients.size(0), -1).norm(2, dim=-1).mean() - else: grad_norm = gradients.norm(2) - loss += torch.log1p((grad_norm - 1) ** 2) - return loss/len(disc_interpolated_output) + net_d.zero_grad() + + gradients = torch.autograd.grad( + outputs=loss_disc, + inputs=interpolated, + retain_graph=True, + allow_unused=True, + materialize_grads=True + )[0] + + if gradients.ndim<=1: gradients = gradients.unsqueeze(0) + grad_norm = gradients.view(gradients.size(0), -1).square().sum(-1).sqrt() + loss = ((grad_norm - 1) ** 2).mean() + return loss # Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/loss.py # LICENSE: https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE @@ -536,7 +553,7 @@ class MultiScaleMelSpectrogramLoss(torch.nn.Module): # Combine losses for this scale scale_losses.append(scale_loss) - total_loss = sum(scale_losses) + total_loss = sum(scale_losses)/len(scale_losses) # Adjust fmin/fmax based on the losses if self.adjustment_factor>0: self.adjust_fmin_fmax([loss.item() for loss in scale_losses]) @@ -556,16 +573,12 @@ def discriminator_loss( disc_generated_outputs: List[torch.Tensor] ): loss = 0 - # r_losses = [] - # g_losses = [] disc_losses = [] for dr, dg in zip(disc_real_outputs, disc_generated_outputs): r_loss = torch.mean((1 - dr) ** 2) g_loss = torch.mean(dg**2) L = r_loss + g_loss loss += L - # r_losses.append(r_loss.item()) - # g_losses.append(g_loss.item()) disc_losses.append(L) return loss, disc_losses diff --git a/training_cli.py b/training_cli.py index 9afe434..1403717 100644 --- a/training_cli.py +++ b/training_cli.py @@ -699,8 +699,8 @@ def train_and_evaluate( balancer_g.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema balancer_d.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema - if total_loss>>[lowest loss]: {least_loss:.3f}<<<") if hps.save_best_model: