implemented loss balancer

custom_nodes/rvc_nodes.py

@@ -246,7 +246,7 @@ class RVCProcessDatasetNode:
         assert dataset, "Please upload a dataset!"
         
         f0_method = pitch_extraction_params.get("f0_method", "")
-        cached_params = [model_name, dataset, period, overlap, max_volume, mute_ratio, sr, f0_method, audio_processor]
+        cached_params = [dataset, period, overlap, max_volume, mute_ratio, sr, f0_method, audio_processor]
         crepe_hop_length = pitch_extraction_params.get("crepe_hop_length",160)
 
         if "crepe" in f0_method: cached_params.append(crepe_hop_length)
@@ -367,6 +367,7 @@ class RVCTrainParamsNode:
                 c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
                 c_sts=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
                 c_gp=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
+                use_balancer=("BOOLEAN",dict(default=False))
             )
         }
 
@@ -377,8 +378,7 @@ class RVCTrainParamsNode:
 
     CATEGORY = CATEGORY
 
-    def init(self,batch_size=4,c_mel=45,c_kl=1.,c_fm=1.,c_mfcc=0.,c_lfcc=0.,c_hd=0.,c_sts=0.,c_gp=0.,):
-        return (dict(batch_size=batch_size,c_mel=c_mel,c_kl=c_kl,c_fm=c_fm,c_mfcc=c_mfcc,c_gp=c_gp,c_lfcc=c_lfcc,c_hd=c_hd,c_sts=c_sts),)
+    def init(self,**kwargs): return (kwargs,)
 
 class RVCTrainModelNode:
     
@@ -407,7 +407,7 @@ class RVCTrainModelNode:
                 train_index=("BOOLEAN",{"default": True}),
                 retrain=("BOOLEAN",{"default": False}),
                 save_best_model=("BOOLEAN",{"default": True}),
-                best_model_threshold=("INT",dict(default=50,min=20,max=80)),
+                best_model_threshold=("INT",dict(default=30,min=10,max=50)),
                 log_every_epoch=("FLOAT",dict(default=1.,min=0.,max=2.,step=.1)),
                 num_workers=("INT",dict(default=1,min=1,max=16))
             )
@@ -443,7 +443,7 @@ class RVCTrainModelNode:
         sample_rate = rvc_dataset_pipe["sample_rate"]
         name = rvc_dataset_pipe["name"]
         dataset_dir = rvc_dataset_pipe["dataset_dir"]
-        cache_name = get_hash(pretrained_G,pretrained_D,*rvc_training_params.items())
+        cache_name = get_hash(dataset_dir,pretrained_G,pretrained_D,*rvc_training_params.items())
         model_dir = os.path.join(output_path,"logs",cache_name)
         if_f0 = rvc_dataset_pipe["if_f0"]
 
@@ -470,8 +470,8 @@ class RVCTrainModelNode:
         hparams.save_best_model = save_best_model
         hparams.best_model_threshold = best_model_threshold
         hparams.log_every_epoch = log_every_epoch
-        hparams.train.num_workers = num_workers
         hparams.train.update(rvc_training_params)
+        hparams.train.num_workers = num_workers
 
         file_index = self.train_index(dataset_dir, sample_rate, name) if train_index else None
         model_path = os.path.join(BASE_MODELS_DIR,"RVC",f"{name}_{sample_rate}.pth")

lib/infer_pack/commons.py

@@ -2,20 +2,71 @@ import math
 import torch
 from torch.nn import functional as F
 
-def median_pool1d(signal, kernel_size, dim=-1,stride=1):
+def median_pool1d(signal, kernel_size, dim=-1, stride=1):
+    """
+    Apply a 1D median filter along a specified dimension for a 3D tensor.
+
+    Args:
+        signal (torch.Tensor): The input tensor of shape (batch_size, F, T).
+        kernel_size (int): The size of the median filter kernel.
+        dim (int): The dimension along which to apply the median filter.
+        stride (int): The stride of the sliding window.
+
+    Returns:
+        torch.Tensor: The filtered tensor.
+    """
+    # Ensure the input tensor is 3D (batch_size, F, T)
+    if signal.dim() != 3:
+        raise ValueError("Input tensor must be 3D (batch_size, F, T)")
+
+    # Adjust kernel size if larger than input dimension
+    input_size = signal.size(dim)
+    if kernel_size > input_size:
+        kernel_size = input_size
+
     # Pad the input to handle edge cases
     padding = (kernel_size - 1) // 2
+    signal = F.pad(signal, (padding, padding), mode='reflect')
+    if dim == -1:
+        signal = F.pad(signal, (padding, padding), mode='reflect')
+    elif dim == -2:
+        signal = F.pad(signal, (0, 0, padding, padding), mode="reflect")
 
-    input_padded = F.pad(signal, (padding, padding), mode='reflect')
-    
     # Unfold the input to create sliding windows
-    unfolded = input_padded.unfold(-1, kernel_size, stride)
-    
+    unfolded = signal.unfold(dim, kernel_size, stride)
+
     # Compute the median along the last dimension
-    median = unfolded.median(dim=dim).values
-    
+    median = unfolded.median(dim=-1).values
+
     return median
 
+def autocorrelation1d(log_magnitude, max_lag=None):
+    """
+    Computes the autocorrelation of log-magnitude spectrogram.
+    :param log_magnitude: Log magnitude spectrogram (batch, channels, time)
+    :param max_lag: Maximum lag for which to compute autocorrelation.
+    :return: Autocorrelation of the log-magnitude spectrogram.
+    """
+    if max_lag is None:
+        max_lag = log_magnitude.size(-1) // 2  # Default to half the signal length
+
+    # Ensure the signal is zero-mean
+    log_magnitude = log_magnitude - log_magnitude.mean(dim=-1, keepdim=True)
+    
+    # Compute autocorrelation using FFT-based convolution
+    padded_log_magnitude = F.pad(log_magnitude, (max_lag, 0), mode='reflect')
+    
+    # Ensure the kernel has the correct shape for conv1d
+    kernel = log_magnitude
+    padded_log_magnitude = padded_log_magnitude
+    
+    autocorr = F.conv1d(padded_log_magnitude, kernel, groups=log_magnitude.size(1))
+    
+    # Take only the relevant part
+    autocorr = F.normalize(autocorr[:,:max_lag + 1], p=2, dim=-1)
+    
+    return autocorr.nan_to_num(0)
+
 def init_weights(m, mean=0.0, std=0.01):
     classname = m.__class__.__name__
     if classname.find("Conv") != -1:
@@ -179,3 +230,7 @@ def clip_grad_value_(parameters, clip_value, norm_type=2, batch_size=1):
             p.grad.data.clamp_(min=-clip_value, max=clip_value)
     total_norm = total_norm ** (1.0 / norm_type)
     return total_norm
+
+def serialize_tensor(tensor):
+    if hasattr(tensor,"item"): return tensor.item()
+    else: return float(f"{float(tensor):.3f}")

lib/train/losses.py

@@ -1,30 +1,179 @@
 import torch
 import torchaudio
 import torch.nn.functional as F
-from ..infer_pack.commons import median_pool1d
 
-def log1p_autocorrelation(log_magnitude, max_lag=None):
-    """
-    Computes the autocorrelation of log-magnitude spectrogram.
-    :param log_magnitude: Log magnitude spectrogram (batch, channels, time)
-    :param max_lag: Maximum lag for which to compute autocorrelation.
-    :return: Autocorrelation of the log-magnitude spectrogram.
-    """
-    if max_lag is None:
-        max_lag = log_magnitude.size(-1) // 2  # Default to half the signal length
+from ..utils import gc_collect
+from ..infer_pack.commons import median_pool1d, autocorrelation1d
 
-    # Ensure the signal is zero-mean
-    log_magnitude = log_magnitude - log_magnitude.mean(dim=-1, keepdim=True)
+class LossBalancer:
+    model: torch.nn.Module
+
+    def __init__(self, model, initial_weights={}, historical_losses={}, ema_weights={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_partial=False, use_pareto=True, use_norm=False):
+        """
+        Initializes the LossBalancer with optional initial weights, partial gradient computation, and EMA.
+        
+        Args:
+            model (torch.nn.Module): The model for which gradients will be computed.
+            initial_weights (dict): Initial weights for each loss term.
+            epsilon (float): A small value to avoid division by zero.
+            partial (bool): Whether to compute gradients only for the output layer.
+        """
+        self.model = model
+        self.epsilon = epsilon
+        self.use_partial = use_partial
+        self.weights_decay = weights_decay
+        self.loss_decay = loss_decay
+        self.initial_weights = initial_weights
+        self.ema_weights = ema_weights
+        self.historical_losses = historical_losses
+        self.active = active
+        self.use_pareto = use_pareto # apply the 20/80 rule to loss balancing
+        self.use_norm = use_norm
+
+
+    def to_dict(self):
+        return dict(
+            epsilon = self.epsilon,
+            use_partial = self.use_partial,
+            weights_decay = self.weights_decay,
+            loss_decay = self.loss_decay,
+            ema_weights = self.ema_weights,
+            initial_weights = self.initial_weights,
+            historical_losses = self.historical_losses,
+            active = self.active,
+            use_pareto = self.use_pareto, # apply the 20/80 rule to loss balancing
+            use_norm = self.use_norm
+        )
+
+    def update_ema_weights(self, new_weights):
+        """Updates the EMA of the weights."""
+        if not self.ema_weights: self.ema_weights = new_weights
+        else: self.ema_weights = {
+                k: self.weights_decay * self.ema_weights.get(k, 0.) + (1 - self.weights_decay) * new_weights[k]
+                for k in new_weights
+            }
     
-    # Compute autocorrelation using FFT-based convolution
-    padded_log_magnitude = F.pad(log_magnitude, (0, max_lag))
-    # autocorr = torch.corrcoef(padded_log_magnitude)
-    autocorr = F.conv1d(padded_log_magnitude, log_magnitude.flip(dims=[-1]))
+    def update_historical_losses(self, new_losses):
+        """Updates the EMA of the weights."""
+        if not self.historical_losses: self.historical_losses = new_losses
+        else: self.historical_losses = {
+                k: self.loss_decay * self.historical_losses.get(k, 0.) + (1 - self.loss_decay) * new_losses[k]
+                for k in new_losses
+            }
+
+    def calculate_loss_slope(self, key, current_loss):
+        """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 = abs((current_loss-ema_loss)/ema_loss)  # relative loss change
+        if slope>1: slope*=slope # punish large changes
+        return slope
     
-    # Take only the relevant part and normalize
-    autocorr = F.normalize(autocorr[:, :max_lag + 1],p=2,dim=-1)
-    
-    return autocorr.nan_to_num(0)
+    def calculate_gradients(self, current_loss):
+        """Calculates the gradient norm of the current loss wrt the model params."""
+        self.model.zero_grad()  # Clear previous gradients
+
+        # Backward pass with partial or full gradient computation
+        model_params = list(self.model.parameters())[-1] if self.use_partial else self.model.parameters()
+        output_params = torch.autograd.grad([current_loss], [model_params], retain_graph=True, only_inputs=True)
+
+        # Compute L2 gradient norm
+        grad_norm = torch.stack([param.norm() for param in output_params if hasattr(param, "norm")]).norm(2)
+        del output_params
+        return grad_norm
+
+    def pareto_normalizer(self, data: dict):
+        # Sort the data in descending order based on the values
+        sorted_data = sorted(data.items(), key=lambda x: x[1], reverse=True)
+        
+        # Calculate the number of elements in the top 20%
+        top_20_count = max(int(0.2 * len(sorted_data)),1)
+        
+        # Calculate the weights
+        weights = [0] * len(sorted_data)
+        total_weight = 1.0
+        top_20_weight = 0.8 * total_weight
+        remaining_weight = total_weight - top_20_weight
+        
+        # Assign weights to the top 20%
+        for i in range(top_20_count):
+            weights[i] = top_20_weight / top_20_count
+        
+        # Assign weights to the remaining 80%
+        for i in range(top_20_count, len(sorted_data)):
+            weights[i] = remaining_weight / (len(sorted_data) - top_20_count)
+        
+        # Normalize the weights
+        total_weight = sum(weights)
+        normalized_weights = [w / total_weight for w in weights]
+        
+        # Combine the keys with their normalized weights
+        normalized_data = {sorted_data[i][0]: normalized_weights[i] for i in range(len(sorted_data))}
+        
+        return normalized_data
+
+    def on_train_batch_start(self, losses: dict):
+        """
+        Balances the loss terms based on their gradients before each training batch.
+        
+        Args:
+            losses (dict): Dictionary of individual loss terms (scalar tensors).
+        
+        Returns:
+            balanced_loss (torch.Tensor): Weighted sum of losses.
+            new_weights (dict): Updated weights for each loss term.
+        """
+        
+        if not self.initial_weights: self.initial_weights = {k: 1.0 for k in losses} # initialize weights
+        if not self.active: return sum(v*self.initial_weights[k] if k in self.initial_weights else v for k,v in losses)
+
+        gradients = {}
+        valid_losses = {}
+
+        # Process each loss, skip if the corresponding weight is 0
+        for key, loss in losses.items():
+            weight = self.initial_weights[key] if key in self.initial_weights else 1.
+            if weight == 0 or loss == 0 or not isinstance(loss,torch.Tensor): continue  # Skip loss with weight 0 or constants
+
+            if self.use_norm:
+                # Compute L2 gradient norm
+                grad_norm = self.calculate_gradients(loss)
+                gradients[key] = grad_norm.item()
+            else:
+                # Use loss plateau detection with EMA
+                loss_slope = self.calculate_loss_slope(key, loss)
+                gradients[key] = loss_slope.item()  # Smaller slope -> plateau -> lower priority
+            valid_losses[key] = loss
+
+        if not valid_losses: return torch.tensor(0.0)  # If all losses are skipped
+
+        # Calculate loss weights based on gradient magnitudes
+        if self.use_pareto: #20% hardest tasks gets 80% weight
+            normalized_weights = self.pareto_normalizer(gradients)
+        else:
+            total_gradient = sum(gradients.values()) + self.epsilon
+            normalized_weights = {k: w/total_gradient for k, w in gradients.items()}
+
+        # Update EMA weights
+        self.update_ema_weights(normalized_weights)
+        balanced_loss = sum(self.ema_weights[k] * loss for k, loss in valid_losses.items())
+        
+        # update historical losses
+        self.update_historical_losses({k: v.item() for k,v in valid_losses.items()})
+        
+        return balanced_loss
+
+    def on_epoch_end(self, weights_decay=None, loss_decay=None):
+        """
+        Optional method to call at the end of each epoch for logging or further adjustments.
+        """
+        if weights_decay is not None: self.weights_decay=weights_decay
+        if loss_decay is not None: self.loss_decay=loss_decay
+        print(f"|| EMA loss weights: {self.ema_weights} ({self.weights_decay=})")
+        print(f"|| EMA losses: {self.historical_losses} ({self.loss_decay=})")
+        weighted_loss = sum(v*self.ema_weights.get(k,0.) for k,v in self.historical_losses.items())
+        print(f"===> EMA loss: {weighted_loss:.3f} <====")
+        gc_collect()
+
 
 def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=None):
     """
@@ -35,79 +184,78 @@ def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=
     :return: Loss value.
     """
     # Compute autocorrelations
-    original_autocorr = log1p_autocorrelation(original_log_magnitude, max_lag)
-    generated_autocorr = log1p_autocorrelation(generated_log_magnitude, max_lag)
+    original_autocorr = autocorrelation1d(original_log_magnitude, max_lag)
+    generated_autocorr = autocorrelation1d(generated_log_magnitude, max_lag)
 
     # Compute the loss between autocorrelations
     loss = F.smooth_l1_loss(generated_autocorr, original_autocorr)
     
     return loss
-    
-def compute_temporal_envelope(log_magnitude: torch.Tensor, kernel_size=31):
-    """Compute the temporal envelope loss using median pooling and second-order differences."""
-    
-    # Apply median pooling to the log-magnitude spectrogram
-    pooled_stft = median_pool1d(log_magnitude, kernel_size, stride=kernel_size // 2 + 1)
-    
-    # Compute second-order differences on the median-pooled audio
-    diff = torch.diff(pooled_stft, dim=-1, n=2)
-    
-    return diff.nan_to_num(0)
 
-def compute_cepstrals(log_magnitude: torch.Tensor, kernel_size=31):
+def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3):
 
-    pooled_magnitude = median_pool1d(log_magnitude, kernel_size, stride=kernel_size//2+1)
+    # average over axis
+    pooled_magnitude = log_magnitude.mean(dim)
 
-    # Compute the cepstrum
-    cepstral_coefficients = torch.fft.ifft(pooled_magnitude, dim=-2).abs()
+    # use max pool to get the peaks
+    max_filtered_tensor = F.max_pool1d(pooled_magnitude, kernel_size=kernel_size, stride=1, padding=kernel_size//2)
     
-    return cepstral_coefficients.sum(dim=-2).nan_to_num(0)
+    return max_filtered_tensor.nan_to_num(0).sum(-2)
 
-def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor, kernel_size=31):
+def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor,eps=1e-6):
     
     # B?N?T
-    org_mag = torch.log1p(org_stft.abs())
-    gen_mag = torch.log1p(gen_stft.abs())
-    org_temp = torch.log1p(org_stft.abs().sum(-2,keepdim=True))
-    gen_temp = torch.log1p(gen_stft.abs().sum(-2,keepdim=True))
+    org_mag = torch.log10(org_stft.abs()+eps)
+    gen_mag = torch.log10(gen_stft.abs()+eps)
 
     # temporal invariant phase
-    phase_loss = compute_phase_loss(org_temp,gen_temp)
+    phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True))
 
     # temporal envelope
-    gen_temp = compute_temporal_envelope(gen_mag, kernel_size=kernel_size)
-    org_temp = compute_temporal_envelope(org_mag, kernel_size=kernel_size)
+    gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
+    org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
     temporal_loss = F.smooth_l1_loss(gen_temp, org_temp)
 
     # spectral
-    gen_spec = compute_cepstrals(gen_mag,kernel_size=kernel_size)
-    org_spec = compute_cepstrals(org_mag,kernel_size=kernel_size)
+    gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
+    org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
     spectral_loss = F.smooth_l1_loss(gen_spec, org_spec)
 
     return temporal_loss, spectral_loss, phase_loss
 
-def compute_harmonics(stft_matrix: torch.Tensor, kernel_size=31):
+def compute_harmonics(stft_matrix: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-6):
     # Calculate log-magnitude spectrograms
-    mag = torch.log1p(stft_matrix.abs())
-    
-    # Apply median filtering to separate harmonics and percussives
-    harmonic = median_pool1d(mag, kernel_size, stride=kernel_size//2+1)
-    percussive = median_pool1d(mag.transpose(-2,-1), kernel_size, stride=kernel_size//2+1).transpose(-2,-1)
-    
-    return harmonic.nan_to_num(0), percussive.nan_to_num(0)
+    mag = torch.log10(stft_matrix.abs()+eps)
+
+    harmonic_list = []
+    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))
+
+    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))
+
+    # Concatenate the results
+    harmonic = torch.cat(harmonic_list, dim=-1)
+    percussive = torch.cat(percussive_list, dim=-1)
+
+    return harmonic, percussive
 
 def combined_aux_loss(
         original_audio: torch.Tensor, generated_audio: torch.Tensor, sample_rate: int,
         c_mfcc=1., c_lfcc=1., c_hd=1., c_sts=1.,
         n_fft=1024, hop_length=320, win_length=1024,
         n_mfcc=13, n_lfcc=13, n_filter=80,
-        f_min=0., f_max=None):
+        f_min=0., f_max=None, eps=None):
 
     kernel_size = n_fft//hop_length+1
-    if kernel_size % 2 == 0: kernel_size+=1
+    if kernel_size % 2 == 0: kernel_size+=1 #enforce odd kernel size
+    if eps is None: eps = torch.finfo(original_audio.dtype).eps
 
     # Compute STFT once
-    if c_hd+c_sts>0:
+    if c_hd+c_sts+c_mfcc+c_lfcc>0:
+        
         hann_window = torch.hann_window(win_length).to(
             dtype=original_audio.dtype, device=original_audio.device
         )
@@ -129,11 +277,15 @@ def combined_aux_loss(
             return_complex=True,
             onesided=True,
             center=False)
+        
+        if c_sts+c_mfcc+c_lfcc>0:
+            org_mag = torch.log10(original_stft.abs()+eps)
+            gen_mag = torch.log10(generated_stft.abs()+eps)
     
     # Harmonic Loss
     if c_hd>0:
-        original_harmonics, original_percussives = compute_harmonics(original_stft,kernel_size)
-        generated_harmonics, generated_percussives = compute_harmonics(generated_stft,kernel_size)
+        original_harmonics, original_percussives = compute_harmonics(original_stft,eps=eps)
+        generated_harmonics, generated_percussives = compute_harmonics(generated_stft,eps=eps)
         # Define loss terms
         harmonic_loss = F.smooth_l1_loss(generated_harmonics, original_harmonics)
         harmonic_loss += F.smooth_l1_loss(generated_percussives, original_percussives)
@@ -143,37 +295,48 @@ def combined_aux_loss(
 
     # Spectral Temporal Smoothness Loss
     if c_sts>0:
-        temporal_loss, spectral_loss, phase_loss = compute_sts_loss(generated_stft,original_stft,kernel_size)
-        temporal_loss *= c_sts
-        spectral_loss *= c_sts
-        phase_loss *= c_sts
-    else: spectral_loss = temporal_loss = phase_loss = 0
+        # temporal invariant phase
+        phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True)) * c_sts
+        
+    else: phase_loss = 0
 
     # MFCC Loss
     if c_mfcc>0:
-        mfcc_transform = torchaudio.transforms.MFCC(
-            sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
-            melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
-        ).to(original_audio.device)
-        original_mfcc = mfcc_transform(original_audio)
-        generated_mfcc = mfcc_transform(generated_audio)
-        mfcc_loss = F.smooth_l1_loss(generated_mfcc, original_mfcc) * c_mfcc
-    else: mfcc_loss = 0
+        # temporal envelope
+        gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
+        org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
+        temporal_loss = F.smooth_l1_loss(gen_temp, org_temp) * c_mfcc
+
+        # mfcc_transform = torchaudio.transforms.MFCC(
+        #     sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
+        #     melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
+        # ).to(original_audio.device)
+        # original_mfcc = mfcc_transform(original_audio)
+        # generated_mfcc = mfcc_transform(generated_audio)
+        # mfcc_loss = F.smooth_l1_loss(generated_mfcc, original_mfcc) * c_mfcc
+        
+    else: temporal_loss = 0
 
     # LFCC Loss
     if c_lfcc>0:
-        lfcc_transform = torchaudio.transforms.LFCC(
-            sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
-            speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
-        ).to(original_audio.device)
-        original_lfcc = lfcc_transform(original_audio)
-        generated_lfcc = lfcc_transform(generated_audio)
-        lfcc_loss = F.smooth_l1_loss(generated_lfcc, original_lfcc) * c_lfcc
-    else: lfcc_loss = 0
+        # spectral
+        gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
+        org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
+        spectral_loss = F.smooth_l1_loss(gen_spec, org_spec) * c_lfcc
+
+        # lfcc_transform = torchaudio.transforms.LFCC(
+        #     sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
+        #     speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
+        # ).to(original_audio.device)
+        # original_lfcc = lfcc_transform(original_audio)
+        # generated_lfcc = lfcc_transform(generated_audio)
+        # lfcc_loss = F.smooth_l1_loss(generated_lfcc, original_lfcc) * c_lfcc
+    else: spectral_loss = 0
+    mfcc_loss=lfcc_loss=0
     
     return harmonic_loss, temporal_loss, spectral_loss, mfcc_loss, lfcc_loss, phase_loss
 
-def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d):
+def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module):
     loss=0
     # Compute the gradient penalty
     # Randomly interpolate between real and generated data
@@ -186,15 +349,16 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens
     _, disc_interpolated_output, _, _ = net_d(original_audio, interpolated)
     # Compute gradients of discriminator output w.r.t. interpolated data
     for output in disc_interpolated_output:
-        gradient = torch.autograd.grad(
+        net_d.zero_grad()
+        gradients = torch.autograd.grad(
             outputs=output,
             inputs=interpolated,
             grad_outputs=torch.ones(output.size(), device=original_audio.device),
-            create_graph=True,
+            retain_graph=True,
             only_inputs=True
-        )[0]
-        gradient = gradient.view(gradient.size(0), -1)
-        loss += torch.log(gradient.norm(2,dim=-1)).abs().mean() # force gradnorm=1
+        )
+        grad_norm = torch.stack([param.norm() for param in gradients if hasattr(param, "norm")]).norm(2)
+        loss += torch.log10(grad_norm).abs().mean() # force gradnorm=1
     return loss/len(disc_interpolated_output)
 
 def feature_loss(fmap_r, fmap_g):

lib/train/utils.py

@@ -105,6 +105,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
 
     iteration = checkpoint_dict["iteration"]
     learning_rate = checkpoint_dict["learning_rate"]
+    kwargs = checkpoint_dict.get("kwargs", {})
     if (
         optimizer is not None and load_opt == 1
     ):  ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch
@@ -113,15 +114,11 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
     #   except:
     #     traceback.print_exc()
     logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))
-    return model, optimizer, learning_rate, iteration
+    return model, optimizer, learning_rate, iteration, kwargs
 
 
-def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
-    logger.info(
-        "Saving model and optimizer state at epoch {} to {}".format(
-            iteration, checkpoint_path
-        )
-    )
+def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path, **kwargs):
+    logger.info(f"Saving model and optimizer state at epoch {iteration} to {checkpoint_path}")
     if hasattr(model, "module"):
         state_dict = model.module.state_dict()
     else:
@@ -132,6 +129,7 @@ def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path)
             "iteration": iteration,
             "optimizer": optimizer.state_dict(),
             "learning_rate": learning_rate,
+            "kwargs": kwargs
         },
         checkpoint_path,
     )

training_cli.py

@@ -30,7 +30,7 @@ from .lib.train.data_utils import (
     TextAudioCollate,
     DistributedBucketSampler,
 )
-from .lib.train.losses import combined_aux_loss, generator_loss, discriminator_loss, feature_loss, gradient_norm_loss, kl_loss
+from .lib.train.losses import LossBalancer, 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):
@@ -235,33 +235,27 @@ def run(rank, n_gpus, hps, device):
         net_d = DDP(net_d)
 
     try:  # resume training
-        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)
+        _, _, _, epoch_str, d_kwargs = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)
         if rank == 0: logger.info("loaded D")
         
-        _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
+        _, _, _, epoch_str, g_kwargs = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
         if rank == 0: logger.info("loaded G")
 
         global_step = (epoch_str - 1) * len(train_loader)
-    except:  # 如果首次不能加载，加载pretrain
-        # traceback.print_exc()
+
+    except Exception as e:
+        logger.error(f"Failed to load saved pretrains: {e}")
         epoch_str = 1
         global_step = 0
         if hps.pretrainG != "":
             if rank == 0:
                 logger.info("loaded pretrained %s" % (hps.pretrainG))
-            print(
-                net_g.module.load_state_dict(
-                    torch.load(hps.pretrainG, map_location="cpu")["model"]
-                )
-            )  ##测试不加载优化器
+            print(net_g.module.load_state_dict(torch.load(hps.pretrainG, map_location="cpu")["model"]))
         if hps.pretrainD != "":
             if rank == 0:
                 logger.info("loaded pretrained %s" % (hps.pretrainD))
-            print(
-                net_d.module.load_state_dict(
-                    torch.load(hps.pretrainD, map_location="cpu")["model"]
-                )
-            )
+            print(net_d.module.load_state_dict(torch.load(hps.pretrainD, map_location="cpu")["model"]))
+        d_kwargs = g_kwargs = {}
 
     scheduler_g = torch.optim.lr_scheduler.ExponentialLR(
         optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str - 2
@@ -272,6 +266,22 @@ def run(rank, n_gpus, hps, device):
 
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
+    try:
+        balancer_state = g_kwargs["balancer"]
+        logger.info(f"Using existing balancer: {balancer_state}")
+        balancer = LossBalancer(net_g,**balancer_state)
+    except Exception as e:
+        logger.error(f"Failed to load balancer state: {e}")
+        balancer = LossBalancer(
+            net_g,
+            use_partial=True,
+            weights_decay=.5 / (1 + np.exp(-10 * (epoch_str / hps.total_epoch - 0.16)))+.5, #sigmoid scaled ema .8 at 20% epoch
+            loss_decay=.8,
+            epsilon=hps.train.eps,
+            active=hps.train.get("use_balancer",False),
+            use_pareto=True,
+            use_norm=hps.train.get("c_gp",0)>0
+            )
     cache = []
     for epoch in range(epoch_str, hps.train.epochs + 1):
         if rank == 0:
@@ -287,6 +297,7 @@ def run(rank, n_gpus, hps, device):
                 logger,
                 [writer, writer_eval],
                 cache,
+                balancer
             )
         else:
             train_and_evaluate(
@@ -301,13 +312,14 @@ def run(rank, n_gpus, hps, device):
                 None,
                 None,
                 cache,
+                balancer
             )
         scheduler_g.step()
         scheduler_d.step()
 
 
 def train_and_evaluate(
-    rank, epoch, hps, nets, optims, _, scaler, loaders, logger, writers, cache
+    rank, epoch, hps, nets, optims, _, scaler, loaders, logger, writers, cache, balancer: "LossBalancer"
 ):
     net_g, net_d = nets
     optim_g, optim_d = optims
@@ -497,25 +509,38 @@ 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",45.)
-                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.get("c_kl",1.)
-                loss_fm = feature_loss(fmap_r, fmap_g) * hps.train.get("c_fm",2.)
+                loss_mel = F.l1_loss(y_mel, y_hat_mel)*hps.train.c_mel if hps.train.get("c_mel",0.)>0 else 0 #default 45
+                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)*hps.train.c_kl if hps.train.get("c_kl",0.)>0 else 0 #default 1
+                loss_fm = feature_loss(fmap_r, fmap_g)*hps.train.c_fm if hps.train.get("c_fm",0.)>0 else 0 #default 2
                 harmonic_loss, temporal_loss, spectral_loss, mfcc_loss, lfcc_loss, phase_loss = combined_aux_loss(
                     wave, y_hat, SR_MAP[hps.sample_rate],
-                    c_lfcc=hps.train.get("c_lfcc",0.),
-                    c_mfcc=hps.train.get("c_mfcc",0.),
-                    c_hd=hps.train.get("c_hd",0.),
-                    c_sts=hps.train.get("c_sts",0.),
+                    c_lfcc=hps.train.c_lfcc if hps.train.get("c_lfcc",0.)>0 else 0.,
+                    c_mfcc=hps.train.c_mfcc if hps.train.get("c_mfcc",0.)>0 else 0.,
+                    c_hd=hps.train.c_hd if hps.train.get("c_hd",0.)>0 else 0.,
+                    c_sts=hps.train.c_sts if hps.train.get("c_sts",0.)>0 else 0.,
                     n_fft=hps.data.filter_length,
                     hop_length=hps.data.hop_length,
                     win_length=hps.data.win_length,
                     n_filter=hps.data.n_mel_channels,
                     f_min=hps.data.mel_fmin,
-                    f_max=hps.data.mel_fmax
+                    f_max=hps.data.mel_fmax,
+                    eps=hps.train.eps
                 )
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
                 aux_loss = harmonic_loss + temporal_loss + spectral_loss + mfcc_loss + lfcc_loss + phase_loss
-                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
+                # loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
+                loss_gen_all = balancer.on_train_batch_start(dict(
+                    loss_gen=loss_gen,
+                    loss_fm=loss_fm,
+                    loss_mel=loss_mel,
+                    loss_kl=loss_kl,
+                    harmonic_loss=harmonic_loss,
+                    temporal_loss=temporal_loss,
+                    spectral_loss=spectral_loss,
+                    phase_loss=phase_loss,
+                    mfcc_loss=mfcc_loss,
+                    lfcc_loss=lfcc_loss
+                    ))
                 
             optim_g.zero_grad()
             scaler.scale(loss_gen_all).backward()
@@ -526,8 +551,8 @@ def train_and_evaluate(
         
         if rank == 0:
             # if epoch>=max(hps.total_epoch*.2,20): #start saving best model after 20% or 20 epochs
-            if loss_gen_all+loss_disc_all<least_loss:
-                least_loss = loss_gen_all+loss_disc_all
+            if loss_gen + loss_fm + loss_mel + loss_kl + aux_loss<least_loss:
+                least_loss = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
                 logger.info(f"[lowest loss] {least_loss=:.3f}: {loss_disc=:.3f} {gradient_penalty=:.3f} <==> {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}")
                 logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {temporal_loss=:.3f}, {spectral_loss=:.3f}, {mfcc_loss=:.3f}, {lfcc_loss=:.3f} {phase_loss=:.3f}")
 
@@ -535,16 +560,37 @@ def train_and_evaluate(
                     if hasattr(net_g, "module"): ckpt = net_g.module.state_dict()
                     else: ckpt = net_g.state_dict()
                     
-                    best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:2.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}"
+                    best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}"
                     status = save_checkpoint(ckpt,best_model_name,epoch,hps)
                     logger.info(f"=== saving best model {best_model_name}: {status=} ===")
                 
                 with open(loss_file,"w") as f:
-                    json.dump(dict(least_loss=least_loss.item(),best_model_name=best_model_name),f)
+                    json.dump(dict(least_loss=least_loss.item(),best_model_name=best_model_name,epoch=epoch,steps=global_step,
+                                   loss_weights = dict(balancer.ema_weights),
+                                   scalar_dict={
+                                       "total/loss/all": commons.serialize_tensor(loss_gen_all+loss_disc_all),
+                                        "total/loss/gen_all": commons.serialize_tensor(loss_gen_all),
+                                        "total/loss/aux": commons.serialize_tensor(aux_loss),
+                                        "total/loss/disc_all": commons.serialize_tensor(loss_disc_all),
+                                        "total/loss/gen": commons.serialize_tensor(loss_gen),
+                                        "total/loss/disc": commons.serialize_tensor(loss_disc),
+                                        "total/loss/fm": commons.serialize_tensor(loss_fm),
+                                        "total/loss/mel": commons.serialize_tensor(loss_mel),
+                                        "total/loss/kl": commons.serialize_tensor(loss_kl),
+                                        "aux/loss/harmonic": commons.serialize_tensor(harmonic_loss),
+                                        "aux/loss/temporal": commons.serialize_tensor(temporal_loss),
+                                        "aux/loss/spectral": commons.serialize_tensor(spectral_loss),
+                                        "aux/loss/mfcc": commons.serialize_tensor(mfcc_loss),
+                                        "aux/loss/lfcc": commons.serialize_tensor(lfcc_loss),
+                                        "aux/loss/phase_loss": commons.serialize_tensor(phase_loss),
+                                        "gradient/grad_norm_disc": commons.serialize_tensor(grad_norm_d),
+                                        "gradient/grad_norm_gen": commons.serialize_tensor(grad_norm_g),
+                                        "gradient/gradient_penalty": commons.serialize_tensor(gradient_penalty),
+                                   }),f,indent=2)
 
             if hps.train.log_interval>0 and global_step % hps.train.log_interval == 0:
                 lr = optim_g.param_groups[0]["lr"]
-                logger.info(f"Train Epoch: {epoch} [{100.0 * epoch / hps.total_epoch:.2f}% complete]")
+                logger.info(f"Train Epoch: {epoch} [{100.0 * (epoch-1) / hps.total_epoch:.2f}% complete]")
                 # Amor For Tensorboard display
                 if loss_mel > 75:
                     loss_mel = 75
@@ -612,6 +658,7 @@ def train_and_evaluate(
             hps.train.learning_rate,
             epoch,
             os.path.join(hps.model_dir, f"G_{saved_epoch}.pth"),
+            balancer=balancer.to_dict()
         )
         utils.save_checkpoint(
             net_d,
@@ -631,6 +678,7 @@ def train_and_evaluate(
         logger.info(f"|| {loss_disc_all=:.3f}: {loss_disc=:.3f}, {gradient_penalty=:.3f}")
         logger.info(f"|| {loss_gen_all=:.3f}: {loss_gen=:.3f}, {loss_fm=:.3f}, {loss_mel=:.3f}, {loss_kl=:.3f}")
         logger.info(f"|| {harmonic_loss=:.3f}, {temporal_loss=:.3f}, {spectral_loss=:.3f}, {mfcc_loss=:.3f}, {lfcc_loss=:.3f}, {phase_loss=:.3f}")
+        balancer.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema
 
     if epoch >= hps.total_epoch and rank == 0:
         logger.info("Training is done. The program is closed.")
@@ -644,7 +692,7 @@ def train_and_evaluate(
                 if os.path.isfile(best_model_path):
                     shutil.copy(best_model_path,os.path.join(
                         os.path.dirname(hps.model_path),
-                        f"{os.path.basename(hps.model_path).split('.')[0]}-best.pth"))
+                        f"{os.path.basename(hps.model_path).split('.')[0]}-lowest.pth"))
 
         status = save_checkpoint(ckpt,hps.name,epoch,hps,model_path=hps.model_path)
         logger.info(f"saving final ckpt {hps.model_path}: {status}")