diff --git a/custom_nodes/rvc_nodes.py b/custom_nodes/rvc_nodes.py index fcb6483..15b62b1 100644 --- a/custom_nodes/rvc_nodes.py +++ b/custom_nodes/rvc_nodes.py @@ -419,7 +419,7 @@ class RVCTrainModelNode: train_index=("BOOLEAN",{"default": True}), retrain=("BOOLEAN",{"default": False}), save_best_model=("BOOLEAN",{"default": True}), - best_model_threshold=("INT",dict(default=30,min=10,max=50)), + best_model_threshold=("INT",dict(default=10,min=1,max=100)), log_every_epoch=("FLOAT",dict(default=1.,min=0.,max=2.,step=.1)), num_workers=("INT",dict(default=1,min=1,max=16)) ) @@ -448,7 +448,7 @@ class RVCTrainModelNode: train_index=True, retrain=False, save_best_model=True, - best_model_threshold=50, + best_model_threshold=10, log_every_epoch=1., num_workers=1): diff --git a/lib/infer_pack/commons.py b/lib/infer_pack/commons.py index 58c5fc1..ae7e31c 100644 --- a/lib/infer_pack/commons.py +++ b/lib/infer_pack/commons.py @@ -71,6 +71,11 @@ def minmax_scale(tensor: torch.Tensor, eps=1e-8): tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min() + eps) return tensor +def compute_correlation(t1: torch.Tensor, t2: torch.Tensor, eps=1e-8): + numerator = (t1 * t2).nan_to_num(0).sum(dim=-1) + denominator = torch.sqrt((t1 ** 2).sum(dim=-1) * (t2 ** 2).sum(dim=-1)+eps).nan_to_num(0) + return (numerator / denominator).nan_to_num(0) + def init_weights(m, mean=0.0, std=0.01): classname = m.__class__.__name__ if classname.find("Conv") != -1: diff --git a/lib/train/losses.py b/lib/train/losses.py index eae0f77..501458c 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -1,14 +1,15 @@ +from typing import Optional import torch import torchaudio.transforms as T import torch.nn.functional as F from ..utils import gc_collect -from ..infer_pack.commons import median_pool1d, minmax_scale +from ..infer_pack.commons import compute_correlation, median_pool1d, minmax_scale 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_pareto=True, use_norm=False): + def __init__(self, model: torch.nn.Module, initial_weights: dict={}, historical_losses: dict={}, ema_weights: dict={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_pareto=True, use_norm=False): """ Initializes the LossBalancer with optional initial weights, partial gradient computation, and EMA. @@ -29,7 +30,6 @@ class LossBalancer: 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, @@ -45,38 +45,41 @@ class LossBalancer: def update_ema_weights(self, new_weights): """Updates the EMA of the weights.""" - if not self.ema_weights: self.ema_weights = new_weights + if not self.ema_weights: self.ema_weights = dict(**new_weights) else: self.ema_weights = { - k: self.weights_decay * self.ema_weights.get(k, 0.) + (1 - self.weights_decay) * new_weights[k] + k: self.weights_decay * self.ema_weights.get(k, 1.) + (1 - self.weights_decay) * new_weights[k] for k in new_weights } + return dict(**self.ema_weights) def update_historical_losses(self, new_losses: dict): """Updates the EMA of the weights.""" - if not self.historical_losses: self.historical_losses = new_losses + 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] - + return dict(**self.historical_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 + slope = (current_loss-ema_loss)/ema_loss # relative loss change return slope - def calculate_gradients(self, key, current_loss): + def calculate_gradients(self, key, current_loss, input): """Calculates the gradient norm of the current loss wrt the model params.""" self.model.zero_grad() # Clear previous gradients # Backward pass to output layer - model_params = list(self.model.parameters())[-1] - output_params = torch.autograd.grad(current_loss, model_params, grad_outputs=torch.ones_like(current_loss, device=model_params.device), retain_graph=True, only_inputs=True)[0] + output_params = torch.autograd.grad(current_loss, [input], retain_graph=True, allow_unused=True, materialize_grads=True)[0] # Compute L2 gradient norm - grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1).mean() - ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon - return grad_norm/ema_loss + # grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1).mean() + if output_params.ndim>1: grad_norm = output_params.view(output_params.size(0), -1).norm(2, dim=-1).mean() + else: grad_norm = output_params.norm(2) + + slope = self.calculate_loss_slope(key, current_loss) + return grad_norm * slope def pareto_normalizer(self, data: dict): # Sort the data in descending order based on the values @@ -104,11 +107,11 @@ class LossBalancer: 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))} + normalized_data = {sorted_data[i][0]: normalized_weights[i]*len(data) for i in range(len(sorted_data))} return normalized_data - def on_train_batch_start(self, losses: dict): + def on_train_batch_start(self, losses: dict, input: Optional[torch.Tensor]=None): """ Balances the loss terms based on their gradients before each training batch. @@ -120,26 +123,32 @@ class LossBalancer: 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.items()) + if len(losses)==0: return 0. # no losses to balance + if not self.initial_weights: self.initial_weights = {k: 1. for k in losses} # initialize weights + if not self.ema_weights: self.ema_weights = {k: 1. for k in losses} # initialize ema weights + if not self.active: + self.update_historical_losses({k: v.item()*self.initial_weights.get(k,1.) for k,v in losses.items() if v>0}) + return sum(v*self.initial_weights.get(k,1.) for k,v in losses.items()) 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. + weight = self.initial_weights.get(key, 1.) if weight == 0 or loss == 0 or not isinstance(loss,torch.Tensor): continue # Skip loss with weight 0 or constants + weighted_loss = loss * weight - if self.use_norm: + if self.use_norm and input is not None: # Compute L2 gradient norm - grad_norm = self.calculate_gradients(key, loss) - gradients[key] = grad_norm.item() + 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, loss) - gradients[key] = loss_slope.item() # Smaller slope -> plateau -> lower priority - valid_losses[key] = loss + loss_slope = self.calculate_loss_slope(key, weighted_loss) + gradients[key] = max(loss_slope.item(), self.epsilon) + valid_losses[key] = weighted_loss if not valid_losses: return torch.tensor(0.0) # If all losses are skipped @@ -148,7 +157,7 @@ class LossBalancer: 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()} + normalized_weights = {k: w/total_gradient*len(gradients) for k, w in gradients.items()} # Update EMA weights self.update_ema_weights(normalized_weights) @@ -167,12 +176,16 @@ class LossBalancer: """ 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} <====") + weights = dict(sorted(self.ema_weights.items(),key=lambda x:x[1],reverse=True)) + losses = dict(sorted(self.historical_losses.items(),key=lambda x:x[1],reverse=True)) + print(f"|| EMA weights: {weights} ({self.weights_decay=})") + print(f"|| EMA losses: {losses} ({self.loss_decay=})") + print(f"===> Weighted EMA loss: {self.weighted_ema_loss:.3f} <====") gc_collect() + @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()) def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitude: torch.Tensor, dim=-1, eps=1e-8): """ @@ -189,14 +202,12 @@ def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitu generated_envelope = minmax_scale(generated_envelope, eps=eps) # Compute the correlation - numerator = (original_envelope * generated_envelope).sum(dim=-1) - denominator = torch.sqrt((original_envelope ** 2).sum(dim=-1) * (generated_envelope ** 2).sum(dim=-1)+eps) - correlation = numerator / denominator + correlation = compute_correlation(original_envelope, generated_envelope, eps=eps) # Compute the loss as negative correlation - loss = 1-correlation.mean() + loss = 1-correlation.abs() - return loss + return loss.mean() def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3, eps=1e-8): """ @@ -261,9 +272,9 @@ def compute_tefs(audio_signal: torch.Tensor, eps=1e-8): envelope = minmax_scale(envelope, eps=eps) # Compute the instantaneous phase from the analytic signal - phase = torch.angle(analytic_signal) + phase = torch.angle(analytic_signal).cos() - return envelope, phase + return envelope.nan_to_num(0), phase.nan_to_num(0) 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 @@ -285,7 +296,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, percussive + return harmonic.nan_to_num(0), percussive.nan_to_num(0) def combined_aux_loss( original_audio: torch.Tensor, generated_audio: torch.Tensor, @@ -333,14 +344,13 @@ def combined_aux_loss( # Define loss terms harmonic_loss = F.l1_loss(generated_harmonics, original_harmonics) harmonic_loss += F.l1_loss(generated_percussives, original_percussives) - harmonic_loss *= c_hd else: harmonic_loss = 0 # temporal invariant phase if c_tsi>0: freq_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-1, eps=eps) temp_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-2, eps=eps) - tsi_loss = (freq_tsi+temp_tsi) * c_tsi + tsi_loss = (freq_tsi+temp_tsi) else: tsi_loss = 0 # Temperol Envelope and Fine Structure Loss @@ -348,10 +358,9 @@ def combined_aux_loss( # temporal envelope gen_te, gen_tfs = compute_tefs(generated_audio, eps=eps) org_te, org_tfs = compute_tefs(original_audio, eps=eps) - phase_diff = torch.remainder(gen_tfs - org_tfs, 2 * torch.pi) + correlation = 1-compute_correlation(gen_tfs,org_tfs, eps=eps).abs() - tefs_loss = F.l1_loss(gen_te, org_te) + phase_diff.sin().abs().mean() #F.l1_loss(gen_tfs, org_tfs) - tefs_loss*=c_tefs + tefs_loss = F.l1_loss(gen_te, org_te) + correlation.mean() else: tefs_loss = 0 return harmonic_loss, tefs_loss, tsi_loss @@ -375,10 +384,12 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens inputs=interpolated, grad_outputs=torch.ones_like(output, device=original_audio.device), retain_graph=True, - only_inputs=True + allow_unused=True, + materialize_grads=True )[0] - grad_norm = gradients.view(gradients.size(0), -1).norm(2, dim=-1)+eps - loss += torch.log1p((grad_norm - 1) ** 2).mean() + 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) def feature_loss(fmap_r, fmap_g): @@ -391,7 +402,6 @@ def feature_loss(fmap_r, fmap_g): return loss - def discriminator_loss(disc_real_outputs, disc_generated_outputs): loss = 0 r_losses = [] diff --git a/training_cli.py b/training_cli.py index 03e47d4..712e564 100644 --- a/training_cli.py +++ b/training_cli.py @@ -295,8 +295,16 @@ def run(rank, n_gpus, hps, device): epsilon=hps.train.eps, active=hps.train.get("use_balancer",False), use_pareto=hps.train.get("use_pareto",False), - use_norm=not hps.train.get("fast_mode",False) - ) + use_norm=not hps.train.get("fast_mode",False), + initial_weights=dict( + loss_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.), + harmonic_loss=hps.train.get("c_hd",0.), + tsi_loss=hps.train.get("c_tsi",0.), + tefs_loss=hps.train.get("c_tefs",0.), + )) cache = [] for epoch in range(epoch_str, hps.train.epochs + 1): train_loader.batch_sampler.set_epoch(epoch) @@ -524,14 +532,14 @@ 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.) + 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.) 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.c_tefs if hps.train.get("c_tefs",0.)>0 else 0., - c_hd=hps.train.c_hd if hps.train.get("c_hd",0.)>0 else 0., - c_tsi=hps.train.c_tsi if hps.train.get("c_tsi",0.)>0 else 0., + c_tefs=hps.train.get("c_tefs",0.), + c_hd=hps.train.get("c_hd",0.), + c_tsi=hps.train.get("c_tsi",0.), n_fft=hps.data.filter_length, hop_length=hps.data.hop_length, win_length=hps.data.win_length, @@ -539,7 +547,6 @@ def train_and_evaluate( ) loss_gen, losses_gen = generator_loss(y_d_hat_g) aux_loss = harmonic_loss + tefs_loss + tsi_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, @@ -548,7 +555,7 @@ def train_and_evaluate( harmonic_loss=harmonic_loss, tsi_loss=tsi_loss, tefs_loss=tefs_loss, - )) + ),input=y_hat) optim_g.zero_grad() scaler.scale(loss_gen_all).backward() @@ -558,44 +565,10 @@ def train_and_evaluate( scaler.update() if rank == 0: - # if epoch>=max(hps.total_epoch*.2,20): #start saving best model after 20% or 20 epochs - if loss_gen + loss_fm + loss_mel + loss_kl + aux_loss {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}") - logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}") - - if hps.save_best_model: - 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:.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,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/tefs": commons.serialize_tensor(tefs_loss), - "aux/loss/tsi": commons.serialize_tensor(tsi_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: + if hps.train.log_interval>0 and global_step % hps.train.log_interval == 0: #tensorboard logging lr = optim_g.param_groups[0]["lr"] 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 @@ -628,7 +601,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()), - "all/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_mel[0]).abs().data.cpu().numpy(), cmap="hot") + "slice/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).abs().data.cpu().numpy(), cmap="hot") } with torch.no_grad(): @@ -677,12 +650,47 @@ def train_and_evaluate( logger.info(f"saving ckpt {save_name}: {status}") if rank == 0: - logger.info(f"====> Epoch {epoch} (Total Loss={(loss_disc_all+loss_gen_all).item():.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}") + total_loss = balancer.weighted_ema_loss + loss_disc_all + logger.info(f"====> Epoch {epoch} ({total_loss=:.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}") 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}, {tefs_loss=:.3f}, {tsi_loss=:.3f}") balancer.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema + if total_loss {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}") + logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}") + + if hps.save_best_model: + 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:.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,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/tefs": commons.serialize_tensor(tefs_loss), + "aux/loss/tsi": commons.serialize_tensor(tsi_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 epoch >= hps.total_epoch and rank == 0: logger.info("Training is done. The program is closed.")